uptimes/uptimes

#! /bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import argparse
import signal
import socket

startTime = 0
uptime = 0
s = 0

socket_file = '/run/uptimesd.sock'
uptimes_db = '/etc/uptimes/uptimes.db'

def main():
    global uptime, s
    parse = argparse.ArgumentParser(description='Uptimes')
    parse.add_argument('-d', '--daemon', action='store_true', help='run as daemon')
    parse.add_argument('-s', '--systemd', action='store_true', help='run as systemd service')


    args = parse.parse_args()
    if args.daemon:
        daemon()
    elif args.systemd:
        daemon(fork=False)
    else:
        # check socket file
        if not os.path.exists(socket_file):
            print('Uptimesd not running')
            sys.exit(1)

        str = comm('get')
        uptime = float(str)

        # format
        years = uptime // 31536000
        uptime = uptime % 31536000
        days = uptime // 86400
        uptime = uptime % 86400
        hours = uptime // 3600
        uptime = uptime % 3600
        minutes = uptime // 60
        uptime = uptime % 60
        seconds = uptime

        print('Total Uptime: %d years %d days %d hours %d minutes %d seconds' % (years, days, hours, minutes, seconds))


def readUpTime():
    global uptime,startTime

    with open('/proc/uptime', 'r') as f:
        uptime = f.readline()
        uptime = uptime.split(' ')[0]
        uptime = float(uptime)
    return uptime

def pathExists():
    # if /etc/uptimes/uptimes.db exist
    if os.path.exists(uptimes_db):
        # check permissions
        if os.stat(uptimes_db).st_mode != 0o600:
            # change permissions
            os.chmod(uptimes_db, 0o600)
        
        # check the content
        with open(uptimes_db, 'r') as f:
            content = f.readline()
            if content == '':
                # write 0
                with open(uptimes_db, 'w') as f:
                    f.write('0')
                    os.chmod(uptimes_db, 0o600)

    else:
        folder = uptimes_db.split('/')
        folder.pop()
        folder = '/'.join(folder)
        if not os.path.exists(folder):
            # create folder with 600
            os.mkdir(folder, 0o600)
        # create uptimes.db with 600
        with open(uptimes_db, 'w') as f:
            f.write('0')
            os.chmod(uptimes_db, 0o600)


def daemon(fork = True):
    global uptime, now, s

    if os.path.exists(socket_file):
        print('Uptimesd already running')
        sys.exit(1)

    # fork
    if fork:
        pid = os.fork()
        if pid > 0:
            sys.exit(0)
        
        os.chdir('/')
        os.setsid()
        os.umask(0)
        # redirect
        sys.stdout.flush()
        sys.stderr.flush()

        si = open('/dev/null', 'r')
        so = open('/dev/null', 'a+')
        se = open('/dev/null', 'a+')

        os.dup2(si.fileno(), sys.stdin.fileno())
        os.dup2(so.fileno(), sys.stdout.fileno())
        os.dup2(se.fileno(), sys.stderr.fileno())
        # fork again
        pid = os.fork()
        if pid > 0:
            sys.exit(0)

        # redirect
        sys.stdout.flush()
        sys.stderr.flush()

        si = open('/dev/null', 'r')
        so = open('/dev/null', 'a+')
        se = open('/dev/null', 'a+')

        os.dup2(si.fileno(), sys.stdin.fileno())
        os.dup2(so.fileno(), sys.stdout.fileno())
        os.dup2(se.fileno(), sys.stderr.fileno())


    # change process name
    import setproctitle
    setproctitle.setproctitle('uptimesd')
    
    s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    s.bind(socket_file)
    # socket file permission 777
    os.chmod(socket_file, 0o777)
    s.listen(1)

    # handle signal
    signal.signal(signal.SIGTERM, sigterm_handler)
    signal.signal(signal.SIGINT, sigterm_handler)
    signal.signal(signal.SIGQUIT, sigterm_handler)
    signal.signal(signal.SIGPIPE, sigterm_handler)

    pathExists()
    # read uptime
    with open(uptimes_db, 'r') as f:
        uptime = f.readline()
        uptime = float(uptime)
    
    while True:
        conn, addr = s.accept()
        data = conn.recv(1024)
        if data == b'':
            conn.close()
            continue

        elif data == b'ping':
            conn.sendall(b'pong')
            conn.close()
            continue

        elif data == b'update':
            updateTime()
            conn.sendall(b'ok')
            conn.close()
            continue

        elif data == b'get':
            updateTime()
            conn.sendall(b'%f' % uptime)
            conn.close()
            continue
        
        else:
            conn.sendall(b'unknown command')
            conn.close()
            continue

def updateTime():
    global uptime, startTime
    now = readUpTime()
    uptime += now - startTime
    pathExists()
    with open(uptimes_db, 'w') as f:
        f.write(str(uptime))

def sigterm_handler(signo, frame):
    # update uptime
    updateTime()
    # close socket
    s.close()
    if os.path.exists(socket_file):
        os.unlink(socket_file)
    # exit
    sys.exit(0)

def comm(command):
    global socket_file, s
    # check socket file
    if not os.path.exists(socket_file):
        print('Uptimesd not running')
        sys.exit(1)
    
    # connect
    s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    s.connect(socket_file)

    # send command
    s.sendall(command.encode('utf-8'))
    data = s.recv(1024)
    return data.decode('utf-8')

if __name__ == '__main__':
    main()