#!/usr/bin/env python3
import argparse
import fcntl
import http.client
import json
import os
import random
import re
import selectors
import signal
import ssl
import stat
import subprocess
import sys
import tempfile
import time
import urllib.error
import urllib.request

COMMAND_TIMEOUT = 15
OUTPUT_LIMIT = 65536
REPORT_LIMIT = 4194304


def fqdn(name):
    name = name.lower()
    if name.endswith('.'):
        name = name[:-1]
    if len(name) > 253 or '.' not in name or any(
        not re.fullmatch(r'[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?', label)
        for label in name.split('.')
    ):
        raise ValueError('invalid fqdn')
    return name


def private_read(path, limit):
    fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW)
    with os.fdopen(fd, 'rb') as handle:
        st = os.fstat(handle.fileno())
        if not stat.S_ISREG(st.st_mode) or st.st_uid != 0 or st.st_mode & 0o077:
            raise ValueError('unsafe permissions: ' + path)
        data = handle.read(limit + 1)
    if len(data) > limit:
        raise ValueError('file too large: ' + path)
    return data


def config(path):
    values = {}
    for line in private_read(path, 8192).decode('utf-8').splitlines():
        if not line.strip() or line.lstrip().startswith('#'):
            continue
        match = re.fullmatch(r'([a-z_]+)=(\S+)', line)
        if not match:
            raise ValueError('invalid config line')
        key, value = match.groups()
        if key in values:
            raise ValueError('duplicate config key')
        values[key] = value
    if values.keys() - {'server', 'fqdn', 'state_dir', 'bootstrap_file'}:
        raise ValueError('unknown config key')
    if not all(values.get(key) for key in ('server', 'fqdn', 'state_dir')):
        raise ValueError('missing server, fqdn or state_dir')
    match = re.fullmatch(r'https://([a-zA-Z0-9.-]+)(?::([0-9]{1,5}))?/?', values['server'])
    if not match:
        raise ValueError('invalid server URL')
    host = fqdn(match[1])
    port = int(match[2] or 443)
    if not 1 <= port <= 65535:
        raise ValueError('invalid port')
    values['server'] = 'https://' + host + ':' + str(port)
    values['fqdn'] = fqdn(values['fqdn'])
    if not os.path.isabs(values['state_dir']):
        raise ValueError('state_dir must be absolute')
    return values


def run_command(argv, timeout=COMMAND_TIMEOUT, limit=OUTPUT_LIMIT):
    try:
        child = subprocess.Popen(argv, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE, start_new_session=True)
    except OSError:
        return {'stdout': '', 'stderr': 'patchouli: command unavailable\n', 'exit_code': 127}
    streams = {'stdout': bytearray(), 'stderr': bytearray()}
    selector = selectors.DefaultSelector()
    for name in streams:
        selector.register(getattr(child, name), selectors.EVENT_READ, name)
    deadline = time.monotonic() + timeout
    failure = None
    try:
        while selector.get_map() or child.poll() is None:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                failure = 'timeout'
                break
            if not selector.get_map():
                try:
                    child.wait(timeout=remaining)
                except subprocess.TimeoutExpired:
                    failure = 'timeout'
                break
            for key, _ in selector.select(min(remaining, 0.1)):
                chunk = os.read(key.fd, 8192)
                if not chunk:
                    selector.unregister(key.fileobj)
                    key.fileobj.close()
                    continue
                output = streams[key.data]
                available = limit - len(output)
                output.extend(chunk[:available])
                if len(chunk) > available:
                    failure = 'output limit; command output truncated'
                    break
            if failure:
                break
    finally:
        if failure:
            try:
                os.killpg(child.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
        child.wait()
        selector.close()
        for name in streams:
            getattr(child, name).close()
    result = {name: data.decode('utf-8', errors='replace') for name, data in streams.items()}
    if failure:
        diagnostic = '\npatchouli: ' + failure + '\n'
        result['stderr'] = result['stderr'][:max(0, limit - len(diagnostic))] + diagnostic[:limit]
    result['exit_code'] = 124 if failure else (child.returncode if child.returncode >= 0 else 128 - child.returncode)
    return result


def collect_commands(os_id):
    commands = [
        ('os-release', ['/bin/cat', '/etc/os-release']),
        ('uname', ['/usr/bin/uname', '-r']),
        ('architecture', ['/usr/bin/uname', '-m']),
    ]
    if os_id in ('debian', 'ubuntu'):
        commands.append(('packages', ['/usr/bin/dpkg-query', '-W', '-f=${Package}\t${Version}\t${Architecture}\t${db:Status-Status}\n']))
    elif os_id in ('rocky', 'almalinux', 'cloudlinux'):
        commands.append(('packages', ['/usr/bin/rpm', '-qa', '--qf', '%{NAME}\t%{EPOCHNUM}\t%{VERSION}\t%{RELEASE}\t%{ARCH}\n']))
    else:
        raise ValueError('unsupported Linux distribution')
    return commands


def collect(c):
    with open('/etc/os-release', encoding='utf-8') as handle:
        release = handle.read(8192)
    match = re.search(r'^ID=[\"\']?([a-z0-9_-]+)[\"\']?$', release, re.MULTILINE)
    if not match:
        raise ValueError('missing Linux distribution ID')
    results = {name: run_command(argv) for name, argv in collect_commands(match[1])}
    return {'fqdn': c['fqdn'], 'os': 'Linux', 'commands': results}


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def encode(data):
    return json.dumps(data, ensure_ascii=True, separators=(',', ':'), sort_keys=True).encode('utf-8')


def post(c, path, key, data):
    def deadline(signum, frame):
        raise TimeoutError('HTTPS deadline exceeded')

    previous_handler = signal.signal(signal.SIGALRM, deadline)
    previous_timer = signal.setitimer(signal.ITIMER_REAL, 30)
    try:
        return https_post(c, path, key, data)
    except TimeoutError:
        raise ValueError('HTTPS transport failed') from None
    finally:
        signal.setitimer(signal.ITIMER_REAL, *previous_timer)
        signal.signal(signal.SIGALRM, previous_handler)


def https_post(c, path, key, data):
    if not isinstance(key, str) or not re.fullmatch(r'[A-Za-z0-9_-]{20,256}', key):
        raise ValueError('invalid credential')
    body = encode(data)
    if len(body) > REPORT_LIMIT:
        raise ValueError('report exceeds 4 MiB limit')
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), NoRedirect(),
        urllib.request.HTTPSHandler(context=ssl.create_default_context()))
    request = urllib.request.Request(c['server'] + path, data=body, method='POST',
        headers={'Authorization': 'Bearer ' + key, 'Content-Type': 'application/json'})
    try:
        try:
            response = opener.open(request, timeout=20)
        except urllib.error.HTTPError as error:
            response = error
        with response:
            payload = response.read(OUTPUT_LIMIT + 1)
            if len(payload) > OUTPUT_LIMIT:
                raise ValueError('HTTPS response exceeds limit')
            return response.code, payload
    except (OSError, urllib.error.URLError, http.client.HTTPException):
        raise ValueError('HTTPS transport failed') from None


def main():
    if os.geteuid() != 0:
        raise ValueError('must run as root')
    os.umask(0o077)
    os.environ.clear()
    os.environ.update(PATH='/bin:/sbin:/usr/bin:/usr/sbin', LC_ALL='C')
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', default='/etc/patchouli-reporter.conf')
    mode = parser.add_mutually_exclusive_group()
    mode.add_argument('--register', action='store_true')
    mode.add_argument('--collect-only', action='store_true')
    parser.add_argument('--jitter', action='store_true')
    args = parser.parse_args()
    c = config(args.config)
    st = os.lstat(c['state_dir'])
    if not stat.S_ISDIR(st.st_mode) or st.st_uid != 0 or stat.S_IMODE(st.st_mode) != 0o700:
        raise ValueError('state_dir must be a root-owned 0700 directory')
    lockfd = os.open(c['state_dir'] + '/lock', os.O_WRONLY | os.O_CREAT | os.O_NOFOLLOW, 0o600)
    with os.fdopen(lockfd, 'w'):
        try:
            fcntl.flock(lockfd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            raise ValueError('reporter already running') from None
        if args.jitter:
            time.sleep(random.randrange(121))
        state = c['state_dir'] + '/credential.json'
        if args.collect_only:
            report = encode(collect(c))
            if len(report) > REPORT_LIMIT:
                raise ValueError('report exceeds 4 MiB limit')
            sys.stdout.buffer.write(report + b'\n')
            return
        if args.register:
            if os.path.lexists(state):
                raise ValueError('local credential exists; registration refused')
            if not c.get('bootstrap_file'):
                raise ValueError('bootstrap_file is required')
            try:
                key = private_read(c['bootstrap_file'], 1024).decode('ascii')
                if key.endswith('\n'):
                    key = key[:-1]
            except UnicodeError:
                raise ValueError('invalid credential') from None
            status, body = post(c, '/v1/register', key, {'fqdn': c['fqdn']})
            if status != 201:
                raise ValueError('registration rejected: HTTP ' + str(status))
            try:
                credential = json.loads(body)
                valid = (isinstance(credential, dict) and credential.get('fqdn') == c['fqdn'] and
                         isinstance(credential.get('key'), str) and
                         re.fullmatch(r'[A-Za-z0-9_-]{20,256}', credential['key']))
            except (ValueError, UnicodeError):
                valid = False
            if not valid:
                raise ValueError('invalid registration response')
            fd, temp = tempfile.mkstemp(prefix='credential.', dir=c['state_dir'])
            try:
                with os.fdopen(fd, 'wb') as handle:
                    handle.write(encode(credential) + b'\n')
                os.replace(temp, state)
            finally:
                if os.path.exists(temp):
                    os.unlink(temp)
            print('registered ' + c['fqdn'])
            return
        credential = json.loads(private_read(state, 4096))
        if not isinstance(credential, dict) or credential.get('fqdn') != c['fqdn']:
            raise ValueError('credential fqdn mismatch')
        status, _ = post(c, '/v1/report', credential.get('key'), collect(c))
        if not 200 <= status < 300:
            raise ValueError('report rejected: HTTP ' + str(status))


if __name__ == '__main__':
    try:
        main()
    except (ValueError, OSError) as error:
        print(str(error), file=sys.stderr)
        sys.exit(1)