"""Bounded, recording TCP egress broker and Linux transparent guest relay.

Guest metadata is untrusted. The outside broker makes all destination decisions.
No TLS interception: encrypted stream bytes remain encrypted evidence.
"""
import argparse
import ipaddress
import json
import os
import pathlib
import select
import socket
import socketserver
import stat
import struct
import threading
import time
import uuid

from evidence import Evidence

MAX_STREAM = 16 * 1024 * 1024
TTL = 600
STOP = pathlib.Path('/var/lib/decoy/STOP')


def public_ip(value):
    ip = ipaddress.IPv4Address(value)
    if not ip.is_global or ip.is_multicast or ip.is_reserved:
        raise ValueError('destination is not public unicast IPv4')
    return str(ip)


def read_message(sock):
    data = bytearray()
    while len(data) < 4096:
        char = sock.recv(1)
        if char == b'\n':
            message = json.loads(data)
            if not isinstance(message, dict):
                raise ValueError('expected object')
            return message
        if not char:
            raise ValueError('incomplete request')
        data.extend(char)
    raise ValueError('request too large')


def send_message(sock, value):
    sock.sendall(json.dumps(value, separators=(',', ':')).encode() + b'\n')


def allowed_peers(path):
    try:
        fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK)
        with os.fdopen(fd, 'rb') as stream:
            info = os.fstat(stream.fileno())
            if not stat.S_ISREG(info.st_mode) or info.st_uid != 0 or info.st_mode & 0o022 or time.time() - info.st_mtime > 3600:
                return set()
            data = stream.read(65537)
        if len(data) > 65536:
            return set()
        peers = json.loads(data)
        if not isinstance(peers, list):
            return set()
        result = set()
        for peer in peers:
            try:
                if isinstance(peer, str):
                    result.add(public_ip(peer))
            except ValueError:
                pass
        return result
    except (OSError, ValueError, TypeError):
        return set()


def pipe(left, right, record=None, stopped=lambda: False):
    """Preserve chunk ordering; record before forwarding, including binary bytes."""
    streams = {left: (right, 'guest_to_remote'), right: (left, 'remote_to_guest')}
    deadline, total, sequence = time.monotonic() + TTL, 0, 0
    while streams:
        if stopped():
            return 'capture_stopped', total
        if time.monotonic() >= deadline:
            return 'lifetime_limit', total
        ready, _, _ = select.select(list(streams), [], [], max(0, min(1, deadline - time.monotonic())))
        for source in ready:
            target, direction = streams[source]
            data = source.recv(min(16384, MAX_STREAM - total + 1))
            if not data:
                target.shutdown(socket.SHUT_WR)
                del streams[source]
                continue
            if total + len(data) > MAX_STREAM:
                if record:
                    record(sequence, direction, data, forwarded=False)
                return 'byte_limit', total
            if record:
                record(sequence, direction, data, forwarded=True)
            sequence += 1
            target.sendall(data)
            total += len(data)
    return 'eof', total


class Server(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True

    def __init__(self, address, handler, **settings):
        self.__dict__.update(settings)
        self.slots = threading.BoundedSemaphore(32)
        self.failed = threading.Event()
        super().__init__(address, handler)

    def process_request(self, request, address):
        if self.failed.is_set() or STOP.exists() or not self.slots.acquire(False):
            request.close()
            return
        try:
            super().process_request(request, address)
        except BaseException:
            self.slots.release()
            raise

    def process_request_thread(self, request, address):
        try:
            super().process_request_thread(request, address)
        finally:
            self.slots.release()


class Broker(socketserver.BaseRequestHandler):
    def fail(self):
        self.server.failed.set()
        STOP.parent.mkdir(parents=True, exist_ok=True)
        STOP.touch(mode=0o600, exist_ok=True)

    def event(self, kind, **data):
        if self.server.failed.is_set() or STOP.exists():
            raise OSError('evidence failed')
        try:
            self.server.evidence.emit(kind, connection_id=self.cid, reported_guest=self.client_address[0], **data)
        except OSError:
            self.fail()
            raise

    def chunk(self, sequence, direction, data, forwarded):
        try:
            blob = self.server.evidence.save(data)
            self.event('network_chunk', sequence=sequence, direction=direction, forwarded=forwarded, **blob)
        except OSError:
            self.fail()
            raise

    def handle(self):
        self.cid = uuid.uuid4().hex
        self.request.settimeout(10)
        connected = False
        try:
            message = read_message(self.request)
            self.event('network_request', untrusted_request=message)
            if message.get('op') == 'dns':
                name = message.get('name')
                if not isinstance(name, str) or len(name) > 253 or not all(c.isascii() and (c.isalnum() or c in '.-') for c in name):
                    raise ValueError('invalid DNS name')
                addresses = set()
                for answer in socket.getaddrinfo(name, None, socket.AF_INET, socket.SOCK_STREAM):
                    try:
                        addresses.add(public_ip(answer[4][0]))
                    except ValueError:
                        pass
                answers = sorted(addresses)[:16]
                self.event('dns_result', name=name, addresses=answers)
                send_message(self.request, {'ok': True, 'addresses': answers})
                return
            if message.get('op') != 'connect':
                raise ValueError('unknown operation')
            ip = public_ip(message.get('ip'))
            port = message.get('port')
            if type(port) is not int or not 1 <= port <= 65535:
                raise ValueError('invalid port')
            if port not in (80, 443) and ip not in allowed_peers(self.server.peers):
                raise ValueError('port requires registered attacker peer')
            # Numeric IP only: DNS cannot rebind between policy and connect.
            with socket.create_connection((ip, port), timeout=10) as remote:
                self.event('network_connected', destination_ip=ip, destination_port=port)
                send_message(self.request, {'ok': True, 'connection_id': self.cid})
                connected = True
                reason, count = pipe(self.request, remote, self.chunk, lambda: self.server.failed.is_set() or STOP.exists())
                self.event('network_end', reason=reason, bytes=count)
        except (OSError, ValueError, TypeError) as exc:
            try:
                self.event('network_closed', reason=str(exc)[:300])
                if not connected:
                    send_message(self.request, {'ok': False, 'error': 'connection denied or ended'})
            except OSError:
                pass


class Relay(socketserver.BaseRequestHandler):
    def handle(self):
        self.request.settimeout(10)
        try:
            original = self.request.getsockopt(socket.SOL_IP, 80, 16)  # Linux SO_ORIGINAL_DST
            port = struct.unpack('!H', original[2:4])[0]
            ip = socket.inet_ntoa(original[4:8])
            with socket.create_connection(self.server.broker, timeout=10) as remote:
                send_message(remote, {'op': 'connect', 'ip': ip, 'port': port})
                if read_message(remote).get('ok'):
                    pipe(self.request, remote)
        except (OSError, ValueError):
            pass


def dns_reply(packet, broker):
    if len(packet) < 12 or len(packet) > 512:
        return None
    ident, flags, questions, _, _, _ = struct.unpack('!6H', packet[:12])
    if questions != 1 or flags & 0xf800:
        return None
    labels, offset = [], 12
    while offset < len(packet):
        size = packet[offset]
        offset += 1
        if not size:
            break
        if size > 63 or offset + size > len(packet):
            return None
        try:
            labels.append(packet[offset:offset + size].decode('ascii'))
        except UnicodeDecodeError:
            return None
        offset += size
    if offset + 4 != len(packet):
        return None
    qtype, qclass = struct.unpack('!HH', packet[offset:])
    addresses, rcode = [], 0
    if qclass == 1 and qtype == 1:
        try:
            with socket.create_connection(broker, timeout=5) as remote:
                send_message(remote, {'op': 'dns', 'name': '.'.join(labels)})
                reply = read_message(remote)
                addresses = [public_ip(a) for a in reply.get('addresses', [])][:16] if reply.get('ok') else []
                if not addresses:
                    rcode = 3
        except (OSError, ValueError, TypeError):
            rcode = 2
    # Empty answer for AAAA/unsupported types avoids leaking around the broker.
    header = struct.pack('!6H', ident, 0x8080 | (flags & 0x100) | rcode, 1, len(addresses), 0, 0)
    answers = b''.join(b'\xc0\x0c' + struct.pack('!HHIH', 1, 1, 30, 4) + socket.inet_aton(ip) for ip in addresses)
    return header + packet[12:] + answers


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--role', choices=('broker', 'relay'), required=True)
    parser.add_argument('--listen', default='127.0.0.1')
    parser.add_argument('--port', type=int)
    parser.add_argument('--broker', default='10.0.2.100')
    parser.add_argument('--broker-port', type=int, default=19090)
    parser.add_argument('--dns-port', type=int, default=53)
    parser.add_argument('--evidence', default='/var/lib/decoy/network')
    parser.add_argument('--quota-mib', type=int, default=8192)
    parser.add_argument('--peers', default='/var/lib/decoy/allowed-peers.json')
    args = parser.parse_args()
    if args.role == 'broker':
        evidence = Evidence(args.evidence, quota=args.quota_mib * 1024 ** 2)
        server = Server((args.listen, args.port or 19090), Broker, evidence=evidence, peers=args.peers)
    else:
        broker = (args.broker, args.broker_port)
        # Bind DNS before serving TCP so a bind failure prevents partial startup.
        dns = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        dns.bind(('127.0.0.1', args.dns_port))
        def serve_dns():
            while True:
                packet, peer = dns.recvfrom(513)
                response = dns_reply(packet, broker)
                if response:
                    dns.sendto(response, peer)
        threading.Thread(target=serve_dns, daemon=True).start()
        server = Server((args.listen, args.port or 15001), Relay, broker=broker)
    with server:
        server.serve_forever()


if __name__ == '__main__':
    main()