uptimes/uptimes

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

import os
import sys
import argparse
import signal
import socket

startTime = 0
totalTime = 0
uptime = 0
s = 0

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

def readTime():
    global uptime,startTime

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

def timeInit():
    global startTime, totalTime
    with open(uptimes_db, 'r') as f:
        totalTime = f.readline()
        totalTime = float(totalTime)

    print("totalTime: %f" % totalTime)
    sys.stdout.flush()

    startTime = readTime()
    print("startTime: %f" % startTime)
    sys.stdout.flush()

def updateTime():
    global uptime, startTime, totalTime
    now = readTime()
    print("Now: %f" % now)
    sys.stdout.flush()
    print("Uptime: %f" % uptime)
    sys.stdout.flush()
    print("startTime: %f" % startTime)
    sys.stdout.flush()
    print("totalTime: %f" % totalTime)
    sys.stdout.flush()
    totalTime += now - startTime
    print("totalTime: %f" % totalTime)
    sys.stdout.flush()
    with open(uptimes_db, 'w') as f:
        f.write(str(totalTime))

def pathExists():
    # if /etc/uptimes/uptimes.db exist
    if os.path.exists(uptimes_db):
        return True
    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 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 daemon(fork = True):
    global uptime, now, s

    print('Uptimesd starting...')
    sys.stdout.flush()

    if os.path.exists(socket_file):
        print('Uptimesd already running')
        sys.stdout.flush()
        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')
    
    print('Creating socket...')
    sys.stdout.flush()

    if os.path.exists(socket_file):
        os.unlink(socket_file)
    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

    print('Reading db...')
    sys.stdout.flush()

    timeInit()
    print('Uptimesd started')
    sys.stdout.flush()
    
    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' % totalTime)
            conn.close()
            continue
        
        else:
            conn.sendall(b'unknown command')
            conn.close()
            continue


def sigterm_handler(signo, frame):
    print('Uptimesd stopping...')
    sys.stdout.flush()
    # update uptime
    updateTime()

    if os.path.exists(socket_file):
        os.unlink(socket_file)
    # exit
    print('Uptimesd stopped. Goodbye.')
    sys.stdout.flush()

    with open(uptimes_db, 'r') as f:
        uptime = f.readline()
        print('Total Uptime: ' + uptime + ' seconds')
        sys.stdout.flush()

    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)
    try:
        # 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')
    except:
        print('Uptimesd not running')
        sys.exit(1)

if __name__ == '__main__':
    main()