citrixpot/test_network_capture.py

"""Offline tests: fixture sockets only; no external connections."""
import json
import os
import pathlib
import socket
import struct
import tempfile
import threading
from unittest.mock import patch
from types import SimpleNamespace

import network_capture as net
from evidence import Evidence


def run():
    for address in ['127.0.0.1', '10.0.0.1', '169.254.169.254', '192.168.1.1', '224.0.0.1', '0.0.0.0', '100.64.0.1', '192.0.2.1', '::1']:
        try:
            net.public_ip(address)
            raise AssertionError(address)
        except ValueError:
            pass
    assert net.public_ip('8.8.8.8') == '8.8.8.8'
    with tempfile.TemporaryDirectory() as tmp:
        net.STOP = pathlib.Path(tmp) / 'STOP'
        peers = pathlib.Path(tmp) / 'peers.json'
        peers.write_text(json.dumps(['10.0.0.1', '8.8.8.8', 'bad', '127.0.0.1', 134744072]))
        peers.chmod(0o600)
        actual_stat = os.stat(peers)
        # Root ownership is mocked only for this unit check on unprivileged Macs.
        with patch.object(net.os, 'fstat', return_value=SimpleNamespace(
                st_mode=actual_stat.st_mode, st_uid=0, st_mtime=actual_stat.st_mtime)):
            assert net.allowed_peers(peers) == {'8.8.8.8'}
        evidence = Evidence(pathlib.Path(tmp) / 'evidence')
        server = net.Server(('127.0.0.1', 0), net.Broker, evidence=evidence, peers=str(pathlib.Path(tmp) / 'absent'))
        worker = threading.Thread(target=server.serve_forever, daemon=True)
        worker.start()
        original_connect = socket.create_connection
        outgoing = []
        def fixture_connect(destination, timeout=None, **kwargs):
            assert destination == ('8.8.8.8', 80), destination
            outgoing.append(destination)
            broker_end, fixture_end = socket.socketpair()
            def echo():
                with fixture_end:
                    while True:
                        data = fixture_end.recv(16384)
                        if not data:
                            break
                        fixture_end.sendall(data)
            threading.Thread(target=echo, daemon=True).start()
            return broker_end
        with patch.object(net.socket, 'create_connection', fixture_connect):
            with original_connect(server.server_address, timeout=3) as client:
                net.send_message(client, {'op': 'connect', 'ip': '127.0.0.1', 'port': 80})
                assert not net.read_message(client)['ok']
            with original_connect(server.server_address, timeout=3) as client:
                net.send_message(client, {'op': 'connect', 'ip': '8.8.8.8', 'port': 4444})
                assert not net.read_message(client)['ok']
            with original_connect(server.server_address, timeout=3) as client:
                net.send_message(client, {'op': 'connect', 'ip': '8.8.8.8', 'port': 80})
                assert net.read_message(client)['ok']
                payload = b'\x00malware-fixture\xff\r\n'
                client.sendall(payload)
                data = b''
                while len(data) < len(payload):
                    data += client.recv(1024)
                assert data == payload
                client.shutdown(socket.SHUT_WR)
                assert client.recv(1) == b''
            assert len(outgoing) == 1
        server.shutdown()
        server.server_close()
        rows = [json.loads(line) for line in (evidence.directory / 'events.jsonl').read_text().splitlines()]
        chunks = [r for r in rows if r['event'] == 'network_chunk']
        assert {r['direction'] for r in chunks} == {'guest_to_remote', 'remote_to_guest'}
        for row in chunks:
            assert (evidence.directory / row['artifact']).read_bytes() == payload
            assert row['time_ns'] and row['connection_id']
        # DNS parser rejects malformed/compressed questions; AAAA returns no answer.
        assert net.dns_reply(b'bad', ('unused', 0)) is None
        question = struct.pack('!6H', 7, 0x100, 1, 0, 0, 0) + b'\x07example\x03com\0' + struct.pack('!HH', 28, 1)
        reply = net.dns_reply(question, ('unused', 0))
        assert struct.unpack('!6H', reply[:12])[3] == 0
        # Capturing failure closes streams before further forward and persists STOP.
        evidence.quota = evidence.used
        failing = net.Server(('127.0.0.1', 0), net.Broker, evidence=evidence, peers='absent')
        threading.Thread(target=failing.serve_forever, daemon=True).start()
        with original_connect(failing.server_address, timeout=3) as client:
            net.send_message(client, {'op': 'connect', 'ip': '8.8.8.8', 'port': 80})
            assert client.recv(1) == b''
        assert net.STOP.exists() and failing.failed.is_set()
        failing.shutdown()
        failing.server_close()
        evidence.close()
    print('network capture offline checks passed')


if __name__ == '__main__':
    run()