anndl: Fixed a couple of bugs.

[utils.git] / logtail

#!/usr/bin/python3

import sys, os, getopt, subprocess, time, select

class host(object):
    def __init__(self, hostname):
        self.hostname = hostname
        self.proc = None
        self.lastconn = 0
        self.outbuf = bytearray()
        self.errbuf = bytearray()

    def gotline(self, buf, line):
        if buf == self.outbuf:
            sys.stdout.buffer.write(line + b"\n")
            sys.stdout.buffer.flush()
        elif buf == self.errbuf:
            sys.stderr.buffer.write(line + b"\n")
            sys.stderr.buffer.flush()

    def handle(self, fp, buf):
        data = fp.read(4096)
        if data == b"":
            if len(self.outbuf) > 0:
                self.gotline(self.outbuf, self.outbuf)
            if len(self.errbuf) > 0:
                self.gotline(self.errbuf, self.errbuf)
            self.proc.wait()
            self.proc = None
            sys.stderr.write("loctail: disconnected from %s\n" % self.hostname)
        buf.extend(data)
        while True:
            p = buf.find(b'\n')
            if p < 0:
                break
            self.gotline(buf, buf[:p])
            buf[:p + 1] = b""

    def handleout(self):
        self.handle(self.proc.stdout, self.outbuf)
    def handleerr(self):
        self.handle(self.proc.stderr, self.errbuf)

    def connect(self):
        self.proc = subprocess.Popen(["ssh", self.hostname, "tail", "-F", "/var/log/syslog"],
                                     bufsize=0, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        self.lastconn = time.time()

    def close(self):
        if self.proc is not None:
            self.proc.terminate()
            self.proc.wait()
            self.proc = None

def min2(arg, *args):
    ret = arg
    for arg in args:
        if ret is None or (arg is not None and arg < ret):
            ret = arg
    return ret

def tailloop(hosts):
    try:
        now = time.time()
        while True:
            rfd = []
            to = None
            for host in hosts:
                if host.proc is not None:
                    rfd.append(host.proc.stdout)
                    rfd.append(host.proc.stderr)
                else:
                    to = min2(to, host.lastconn + 10)
            rfd, wfd, efd = select.select(rfd, [], [], None if to is None else max(to - time.time(), 0))
            now = time.time()
            rfd = set(rfd)
            for host in hosts:
                if host.proc is not None:
                    if host.proc.stdout in rfd:
                        host.handleout()
                    elif host.proc.stderr in rfd:
                        host.handleerr()
                else:
                    if now > host.lastconn + 10:
                        host.connect()
    except KeyboardInterrupt:
        for host in hosts:
            host.close()

def usage(out):
    out.write("usage: logtail [-h] HOST...\n")

opts, args = getopt.getopt(sys.argv[1:], "h")
for o, a in opts:
    if o == "-h":
        usage(sys.stdout)
        sys.exit(0)
if len(args) < 1:
    usage(sys.stderr)
    sys.exit(1)

hosts = [host(arg) for arg in args]
tailloop(hosts)