#!/usr/bin/env python3
"""Disposable Linux guest only. Never run on a workstation or collector."""
import base64
import hashlib
import json
import os
import pathlib
import signal
import stat
import subprocess
import sys
import threading
import time

LOCK = threading.Lock()
LIMIT = 16 * 1024 * 1024
ACTIVE = None
ROOTS = ['/var/tmp', '/tmp', '/root', '/netscaler', '/nsconfig', '/usr/local/bin', '/etc']


def emit(event, content=None, **data):
    if content is not None:
        data['base64'] = base64.b64encode(content).decode()
    with LOCK:
        try:
            print(json.dumps(dict(event=event, guest_time_ns=time.time_ns(), **data)), flush=True)
        except (BrokenPipeError, OSError):
            terminate()
            os._exit(1)


def scan():
    result = {}
    for root in ROOTS:
        for folder, dirs, files in os.walk(root, followlinks=False):
            dirs[:] = dirs[:128]
            for name in files:
                path = os.path.join(folder, name)
                try:
                    s = os.lstat(path)
                    if stat.S_ISREG(s.st_mode):
                        result[path] = (s.st_ino, s.st_mtime_ns, s.st_ctime_ns, s.st_size, s.st_mode)
                except OSError:
                    pass
                if len(result) >= 20000:
                    return result
    return result


def watch(stop, previous):
    # ponytail: polling preserves observed versions; syscall trace records fast deletes/other paths.
    while not stop.wait(.25):
        current = scan()
        for path in previous.keys() - current.keys():
            emit('guest_file_deleted', path=path)
        for path, info in current.items():
            if previous.get(path) == info:
                continue
            try:
                fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK)
                with os.fdopen(fd, 'rb') as f:
                    if not stat.S_ISREG(os.fstat(f.fileno()).st_mode):
                        continue
                    data = f.read(LIMIT + 1)
                emit('guest_file_version', content=data[:LIMIT], path=path, size=info[3],
                     mode=oct(info[4]), truncated=len(data) > LIMIT)
            except OSError as e:
                emit('guest_file_capture_error', path=path, error=str(e))
        previous = current


def terminate(*_):
    if ACTIVE is not None:
        try:
            os.killpg(ACTIVE.pid, signal.SIGKILL)
        except ProcessLookupError:
            pass


def open_beneath(path, flags, mode=0o600):
    # Resolve every directory with NOFOLLOW, including mutable /var/tmp/uploads.
    parts = pathlib.PurePosixPath(path).parts
    if parts[0] != '/' or '..' in parts:
        raise ValueError('invalid path')
    directory = os.open('/', os.O_RDONLY | os.O_DIRECTORY)
    try:
        for part in parts[1:-1]:
            new = os.open(part, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=directory)
            os.close(directory)
            directory = new
        return os.open(parts[-1], flags | os.O_NOFOLLOW | os.O_NONBLOCK, mode, dir_fd=directory)
    finally:
        os.close(directory)


def main():
    global ACTIVE
    marker = pathlib.Path('/etc/decoy-disposable-guest')
    if not marker.is_file() or marker.is_symlink():
        raise SystemExit('Refusing execution outside disposable guest')
    maximum = 4 * ((LIMIT + 2) // 3) + 131072
    line = sys.stdin.buffer.readline(maximum + 1)
    if len(line) > maximum or not line.endswith(b'\n'):
        raise SystemExit('request limit')
    data = json.loads(line)
    if data['action'] == 'read':
        path = pathlib.PurePosixPath(data['path'])
        if not str(path).startswith('/netscaler/ns_gui/') or '..' in path.parts:
            raise SystemExit('invalid webroot path')
        fd = open_beneath(str(path), os.O_RDONLY)
        with os.fdopen(fd, 'rb') as f:
            if not stat.S_ISREG(os.fstat(f.fileno()).st_mode):
                raise SystemExit('not regular file')
            raw = f.read(LIMIT + 1)
            emit('guest_http_file', content=raw[:LIMIT], path=str(path), truncated=len(raw) > LIMIT)
        return
    if data['action'] == 'upload':
        raw = base64.b64decode(data['base64'], validate=True)
        if len(raw) > LIMIT:
            raise SystemExit('upload limit')
        name = '/var/tmp/uploads/' + hashlib.sha256(raw).hexdigest() + '.bin'
        os.makedirs('/var/tmp/uploads', mode=0o700, exist_ok=True)
        # Never truncate attacker-selected symlinks, hardlinks, devices or existing files.
        try:
            fd = open_beneath(name, os.O_WRONLY | os.O_CREAT | os.O_EXCL)
        except FileExistsError:
            emit('guest_upload', path=name, bytes=len(raw), error='path exists; offguest original retained')
            return
        with os.fdopen(fd, 'wb') as f:
            f.write(raw)
        emit('guest_upload', path=name, bytes=len(raw))
        return
    if data['action'] != 'execute':
        raise SystemExit('unknown action')
    command = data['command']
    if not isinstance(command, str) or len(command) > 65536:
        raise SystemExit('command limit')
    for sig in (signal.SIGHUP, signal.SIGTERM, signal.SIGINT):
        signal.signal(sig, lambda *_: (terminate(), sys.exit(1)))
    stop = threading.Event()
    watcher = threading.Thread(target=watch, args=(stop, scan()), daemon=True)
    watcher.start()
    proc = subprocess.Popen(['strace', '-f', '-ttt', '-T', '-s', '4096', '-yy',
                             '/bin/bash', '-c', command], cwd='/var/tmp',
                            stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                            start_new_session=True)
    ACTIVE = proc
    emit('guest_process_start', pid=proc.pid, command=command)
    def drain(stream, event):
        while chunk := stream.read1(32768):
            emit(event, content=chunk)
    readers = [threading.Thread(target=drain, args=(proc.stdout, 'guest_stdout'), daemon=True),
               threading.Thread(target=drain, args=(proc.stderr, 'guest_syscalls_stderr'), daemon=True)]
    for reader in readers:
        reader.start()
    try:
        proc.wait(timeout=600)
    except subprocess.TimeoutExpired:
        emit('guest_execution_limit', seconds=600)
    finally:
        try:
            os.killpg(proc.pid, signal.SIGKILL)
        except ProcessLookupError:
            pass
        proc.wait()
        for reader in readers:
            reader.join(3)
        time.sleep(.5)
        stop.set()
        watcher.join(2)
    emit('guest_process_end', returncode=proc.poll())


if __name__ == '__main__':
    main()