#!/usr/bin/env python3
"""Check already-resolved client addresses against a local Private Relay feed."""

import argparse
from bisect import bisect_right
import csv
import io
import ipaddress
import json
from pathlib import Path

DEMO_RANGES = """192.0.2.0/28,ZZ,,,
192.0.2.8/29,ZZ,,,
2001:db8:1::/48,ZZ,,,
"""
DEMO_EVENTS = """request_id,client_ip,outcome
example-a,192.0.2.10,accepted_lead
example-b,192.0.2.10,rejected_form
example-c,2001:db8:1::5,page_view
example-d,198.51.100.7,rejected_form
example-e,203.0.113.8,accepted_lead
example-f,,page_view
"""
FIELDS = ['request_id', 'client_ip', 'outcome']


def range_index(feed):
    by_version = {4: [], 6: []}
    for row in csv.reader(feed, strict=True):
        if not row or row[0].lstrip().startswith('#'):
            continue
        network = ipaddress.ip_network(row[0].strip())
        by_version[network.version].append((int(network.network_address), int(network.broadcast_address)))

    index = {}
    for version, intervals in by_version.items():
        starts, ends = [], []
        for start, end in sorted(intervals):
            if ends and start <= ends[-1] + 1:
                ends[-1] = max(ends[-1], end)
            else:
                starts.append(start)
                ends.append(end)
        index[version] = (starts, ends)
    if not any(starts for starts, _ in index.values()):
        raise ValueError('No CIDR ranges found in the feed.')
    return index


def classify(value, index):
    if '%' in value:
        return 'unresolved'
    try:
        address = ipaddress.ip_address(value.strip())
    except ValueError:
        return 'unresolved'

    if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped:
        address = address.ipv4_mapped
    starts, ends = index[address.version]
    number = int(address)
    position = bisect_right(starts, number) - 1
    return 'feed_match' if position >= 0 and number <= ends[position] else 'no_feed_match'


def review(feed, events):
    index = range_index(feed)
    reader = csv.DictReader(events, strict=True)
    if reader.fieldnames != FIELDS:
        raise ValueError('Expected CSV columns: ' + ','.join(FIELDS))

    report = []
    for row in reader:
        if None in row or any(value is None for value in row.values()):
            raise ValueError(f'CSV record ending at line {reader.line_num} has the wrong number of fields.')
        report.append({
            'request_id': row['request_id'],
            'outcome': row['outcome'],
            'network_context': classify(row['client_ip'], index),
        })
    return report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--demo', action='store_true', help='Use synthetic events and documentation ranges.')
    parser.add_argument('--ranges', type=Path, help='Local Apple egress range CSV snapshot.')
    parser.add_argument('--events', type=Path, help='Local CSV with already-resolved client addresses.')
    args = parser.parse_args()
    try:
        if args.demo:
            if args.ranges or args.events:
                parser.error('--demo runs by itself.')
            report = review(io.StringIO(DEMO_RANGES), io.StringIO(DEMO_EVENTS))
        else:
            if not (args.ranges and args.events):
                parser.error('Provide --ranges and --events, or use --demo.')
            with args.ranges.open(encoding='utf-8-sig', newline='') as feed, args.events.open(
                encoding='utf-8-sig', newline=''
            ) as events:
                report = review(feed, events)
    except (OSError, ValueError, csv.Error) as error:
        parser.exit(1, f'{error}\n')
    print(json.dumps(report, indent=2))


if __name__ == '__main__':
    main()