#!/usr/bin/env python3
"""Offline, deterministic teaching models. They do not modify host networking."""
import argparse
import heapq
import ipaddress
import json
import math
import struct
from statistics import median


def finite(value):
    """True for ordinary numbers; NaN/Infinity are rejected, huge ints are never converted to float."""
    return not isinstance(value,float) or math.isfinite(value)


def L02(window=65536,rtt=.04):
    if not (finite(window) and window>0 and finite(rtt) and rtt>0):
        raise ValueError('window and RTT must be finite numbers greater than 0')
    speed=100_000_000
    return {'kind':'model','packet_1500B_100Mbps_ms':1500*8/speed*1000,
            'three_hop_ms':3*(1500*8/speed+.002)*1000,
            'bdp_bytes':speed*rtt/8,'window_bound_Mbps':min(speed,window*8/rtt)/1e6,
            'rtt_seconds':rtt,'window_bytes':window,
            'window_comparison_Mbps':{str(w):min(speed,w*8/rtt)/1e6 for w in [16384,65536,500000]},
            'assumption':'W/RTT steady-state approximation; RTT is full measured/modelled round trip, not ARQ last-bit waiting interval.'}


def route_lookup(destination,routes):
    address=ipaddress.ip_address(destination)
    eligible=[(ipaddress.ip_network(prefix),hop) for prefix,hop in routes
              if address.version==ipaddress.ip_network(prefix).version
              and address in ipaddress.ip_network(prefix)]
    if not eligible:return None
    network,hop=max(eligible,key=lambda row:row[0].prefixlen)
    return {'prefix':str(network),'next_hop':hop}


L03_ROUTES=[('192.0.2.64/26','on-link A'),('198.51.100.0/24','192.0.2.65'),
            ('198.51.100.128/25','192.0.2.66'),('0.0.0.0/0','192.0.2.65')]
L03_ADDRESSES=['192.0.2.80','198.51.100.20','198.51.100.140','203.0.113.9']


def L03(extra_routes=(),extra_lookups=()):
    """Chapter 03 table; extra routes are appended and extra destinations looked up too."""
    added=[(str(ipaddress.ip_network(prefix)),hop) for prefix,hop in extra_routes]
    addresses=L03_ADDRESSES+[str(ipaddress.ip_address(a)) for a in extra_lookups
                             if str(ipaddress.ip_address(a)) not in L03_ADDRESSES]
    net=ipaddress.ip_network('198.51.100.160/27')
    result={'kind':'model','subnet':{'network':str(net.network_address),
        'broadcast':str(net.broadcast_address),'netmask':str(net.netmask),
        'usable_count':sum(1 for _ in net.hosts())},
        'routes':{a:route_lookup(a,L03_ROUTES+added) for a in addresses},
        'ipv6':str(ipaddress.ip_address('2001:0db8:0000:0000:0000:0000:0000:0001'))}
    if added:result['added_routes']=[{'prefix':p,'next_hop':h} for p,h in added]
    return result


def dns_name(packet,offset):
    """Decode labels and compression, rejecting loops/out-of-bounds input."""
    labels=[];seen=set();end=None;total=0
    while True:
        if offset in seen or offset>=len(packet):raise ValueError('invalid DNS name pointer')
        seen.add(offset);length=packet[offset];offset+=1
        if length & 0xc0 == 0xc0:
            if offset>=len(packet):raise ValueError('short pointer')
            target=((length & 0x3f)<<8)|packet[offset]
            if end is None:end=offset+1
            offset=target;continue
        if length & 0xc0:raise ValueError('unsupported label form')
        if length==0:return '.'.join(labels),end if end is not None else offset
        if length>63 or offset+length>len(packet):raise ValueError('invalid label')
        total+=length+1
        if total>254:raise ValueError('name too long')
        labels.append(packet[offset:offset+length].decode('ascii'));offset+=length


def L04(ttl=300):
    if not 0<=ttl<2**31:raise ValueError('DNS TTL must be between 0 and 2^31-1')
    qname=b'\x06lesson\x04test\x00';question=qname+struct.pack('!HH',1,1)
    query=struct.pack('!6H',0x1234,0x0100,1,0,0,0)+question
    answer=b'\xc0\x0c'+struct.pack('!HHIH',1,1,ttl,4)+ipaddress.ip_address('192.0.2.80').packed
    reply=struct.pack('!6H',0x1234,0x8180,1,1,0,0)+question+answer
    ident,flags,qd,an,ns,ar=struct.unpack('!6H',reply[:12])
    assert ident==struct.unpack('!H',query[:2])[0] and flags & 0x8000
    name,end=dns_name(reply,12);owner,pos=dns_name(reply,end+4)
    typ,cls,ttl,length=struct.unpack('!HHIH',reply[pos:pos+10])
    return {'kind':'constructed DNS fixture, not captured traffic','query_hex':query.hex(),
        'response_hex':reply.hex(),'id':ident,'flags_hex':f'{flags:04x}',
        'header':{'QR':(flags>>15)&1,'RD':(flags>>8)&1,'RA':(flags>>7)&1,'RCODE':flags&15,
                  'QDCOUNT':qd,'ANCOUNT':an,'NSCOUNT':ns,'ARCOUNT':ar},'question_name':name,
        'answer':{'owner':owner,'type':typ,'class':cls,'ttl':ttl,'rdlength':length,
                  'address':str(ipaddress.ip_address(reply[pos+10:pos+10+length]))}}


def L06(window=3,data_loss=True,ack_loss=True):
    """Selective Repeat round model: immediate per-item ACK except scripted losses."""
    if window<1:raise ValueError('window must be at least 1')
    pending=list(range(5));attempts={n:0 for n in pending};buffer=set();delivered=[];base=0;events=[];round_no=0
    while pending:
        round_no+=1;active=[n for n in pending if n<min(pending)+window]
        for n in active:
            attempts[n]+=1
            if data_loss and n==1 and attempts[n]==1:
                events.append(f'round {round_no}: DATA {n} lost');continue
            duplicate=n in buffer or n in delivered
            if not duplicate:buffer.add(n)
            while base in buffer:
                buffer.remove(base);delivered.append(base);base+=1
            if ack_loss and n==0 and attempts[n]==1:
                events.append(f'round {round_no}: DATA {n} received; ACK lost');continue
            events.append(f'round {round_no}: DATA {n} received'+(' duplicate' if duplicate else '')+'; ACK received')
            pending.remove(n)
        if round_no>20:raise RuntimeError('model failed to converge')
    return {'kind':'discrete-round SR model; one round substitutes for a timer cycle',
        'window':window,'data_loss_rule':data_loss,'ack_loss_rule':ack_loss,'rounds':round_no,
        'events':events,'delivered':delivered,'attempts':attempts,
        'limitations':'No byte sequence wrapping, real time, network socket, cwnd or OS TCP; receiver buffer unlimited; each round represents one timer cycle.'}


def shortest_paths(graph,start):
    distances={start:0};parent={};queue=[(0,start)]
    while queue:
        distance,node=heapq.heappop(queue)
        if distance!=distances[node]:continue
        for neighbour,cost in graph[node].items():
            if cost<0:raise ValueError('Dijkstra needs nonnegative costs')
            candidate=distance+cost
            if candidate<distances.get(neighbour,math.inf):
                distances[neighbour]=candidate;parent[neighbour]=node
                heapq.heappush(queue,(candidate,neighbour))
    return distances,parent


def all_shortest_paths(graph,start,target,best):
    """List every simple path whose cost equals the shortest distance (small teaching graphs only)."""
    found=[]
    def walk(node,cost,path):
        if cost>best:return
        if node==target:
            if cost==best:found.append(path)
            return
        for neighbour,step in sorted(graph[node].items()):
            if neighbour not in path:walk(neighbour,cost+step,path+[neighbour])
    walk(start,0,[start])
    return found


def L07(bc_cost=1):
    if bc_cost<0:raise ValueError('Dijkstra needs nonnegative costs')
    graph={'A':{'B':2,'C':5},'B':{'A':2,'C':bc_cost,'D':4},
           'C':{'A':5,'B':bc_cost,'D':1},'D':{'B':4,'C':1}}
    distances,parent=shortest_paths(graph,'A');paths={}
    for node in graph:
        path=[node]
        while path[-1]!='A':path.append(parent[path[-1]])
        paths[node]=list(reversed(path))
    return {'kind':'static link-state model','graph':graph,'costs':distances,'paths':paths,
            'all_shortest_paths':{node:all_shortest_paths(graph,'A',node,distances[node]) for node in graph},
            'tie_breaking':'paths keeps the first parent found with a strictly smaller cost; all_shortest_paths lists equal-cost alternatives.'}


L08_FILTERS={'EIF':'endpoint-independent filtering','ADF':'address-dependent filtering',
             'APDF':'address-and-port-dependent filtering'}


def L08(filtering='ADF',elapsed=0,lifetime=None):
    # EIM creates one external tuple per internal tuple; filtering decides which remote endpoints may enter.
    if filtering not in L08_FILTERS:raise ValueError('filtering must be EIF, ADF or APDF')
    if not (finite(elapsed) and elapsed>=0) or (lifetime is not None and not (finite(lifetime) and lifetime>0)):
        raise ValueError('elapsed must be a finite number >= 0 and lifetime a finite number > 0')
    internal=('10.0.0.2',51000);external=('203.0.113.5',40000)
    contacted=[('198.51.100.10',9000),('198.51.100.20',9000)]
    probes=[('198.51.100.10',9000),('198.51.100.10',9999),('198.51.100.30',9000)]
    expired=lifetime is not None and elapsed>lifetime
    def allowed(source):
        if expired:return False
        if filtering=='EIF':return True
        if filtering=='ADF':return source[0] in {ip for ip,_ in contacted}
        return source in contacted
    return {'kind':f'UDP NAT model: endpoint-independent mapping, {L08_FILTERS[filtering]}',
       'internal':internal,'external_mapping':None if expired else external,
       'filtering':filtering,'outbound_remote_ips':sorted({ip for ip,_ in contacted}),
       'outbound_remote_endpoints':contacted,
       'elapsed_seconds':elapsed,'state_lifetime_seconds':lifetime,'state_expired':expired,
       'inbound_decisions':[{'source':p,'allowed':allowed(p)} for p in probes],
       'limitations':'No actual packets or TCP behavior; state expiry is one artificial timer for the whole mapping.'}


def L10(path_mtu=1500):
    outer_ipv4=20;udp=8;wg_header=16;tag=16
    inner_v4=path_mtu-outer_ipv4-udp-wg_header-tag
    inner_v6=path_mtu-40-udp-wg_header-tag if path_mtu>=1280 else None
    if inner_v4<=0:raise ValueError('path MTU is smaller than the outer overhead')
    return {'kind':'overhead budget model','path_mtu':path_mtu,
       'outer_ipv4_budget_for_inner_ip':inner_v4,
       'outer_ipv6_budget_for_inner_ip':inner_v6,
       'inner_ipv6_minimum_1280_met':{'outer_ipv4':inner_v4>=1280,'outer_ipv6':inner_v6 is not None and inner_v6>=1280},
       'notes':'WG may pad plaintext; chosen inner MTU must fit actual padded packets and path. No universal 1420 guarantee. Outer IPv6 needs a path MTU of at least 1280.',
       'route_policy':{'203.0.113.9':'physical endpoint exemption','198.51.100.20':'tunnel by explicit policy','192.0.2.80':'local on-link'}}


def L11():
    # Deliberately manufactured numbers: cannot rank real proxy protocols.
    a=[10,11,9,12,10,11,10,12,9,50];b=[13,14,12,13,15,14,13,12,14,16]
    def stats(values):
        ordered=sorted(values)
        return {'n':len(values),'median_ms':median(values),
          'p95_nearest_rank_ms':ordered[math.ceil(.95*len(values))-1],
          'mean_ms':sum(values)/len(values),'max_ms':max(values)}
    return {'kind':'synthetic statistical worksheet','A':stats(a),'B':stats(b),
       'raw_A_ms':a,'raw_B_ms':b,'interpretation':'A lower median, B lower tail in this fixture; n=10 is weak tail evidence.'}


def display(value):
    """Round floats only for printing; the functions keep full precision for verification."""
    if isinstance(value,float):
        rounded=round(value,6)
        return int(rounded) if rounded.is_integer() else rounded
    if isinstance(value,dict):return {k:display(v) for k,v in value.items()}
    if isinstance(value,(list,tuple)):return [display(v) for v in value]
    return value


def route_argument(text):
    prefix,sep,hop=text.partition('=')
    if not sep or not hop:raise argparse.ArgumentTypeError('use PREFIX=NEXT_HOP, e.g. 198.51.100.128/26=192.0.2.67')
    try:ipaddress.ip_network(prefix)
    except ValueError as exc:raise argparse.ArgumentTypeError(str(exc))
    return prefix,hop


def address_argument(text):
    try:return str(ipaddress.ip_address(text))
    except ValueError as exc:raise argparse.ArgumentTypeError(str(exc))


def bounded(kind,minimum,strict):
    def parse(text):
        value=kind(text)
        if not finite(value):raise argparse.ArgumentTypeError('must be a finite number')
        if value<minimum or (strict and value==minimum):
            raise argparse.ArgumentTypeError(f'must be {">" if strict else ">="} {minimum}')
        return value
    parse.__name__=kind.__name__  # argparse names the type in its error message
    return parse


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    labs=parser.add_subparsers(dest='lab',required=True,metavar='LAB')
    p=labs.add_parser('L02',help='delay, BDP and window bound')
    p.add_argument('--window',type=bounded(int,0,True),default=65536,help='bytes allowed in flight (default 65536)')
    p.add_argument('--rtt',type=bounded(float,0,True),default=.04,help='round-trip time in seconds (default 0.04)')
    p=labs.add_parser('L03',help='subnet and longest-prefix match')
    p.add_argument('--route',type=route_argument,action='append',default=[],metavar='PREFIX=NEXT_HOP',help='add a route (repeatable)')
    p.add_argument('--lookup',type=address_argument,action='append',default=[],metavar='ADDRESS',help='also look up this destination (repeatable)')
    p=labs.add_parser('L04',help='DNS message fixture')
    p.add_argument('--ttl',type=int,default=300,help='answer TTL in seconds (default 300)')
    p=labs.add_parser('L06',help='Selective Repeat round model')
    p.add_argument('--window',type=bounded(int,1,False),default=3,help='sender window in data units (default 3)')
    p.add_argument('--no-data-loss',action='store_true',help='do not lose the first transmission of DATA 1')
    p.add_argument('--no-ack-loss',action='store_true',help='do not lose the first ACK for DATA 0')
    p=labs.add_parser('L07',help='Dijkstra on the chapter 10 graph')
    p.add_argument('--bc-cost',type=bounded(int,0,False),default=1,help='B-C link cost in both directions (default 1)')
    p=labs.add_parser('L08',help='NAT mapping and filtering')
    p.add_argument('--filter',choices=sorted(L08_FILTERS),default='ADF',help='filtering behavior (default ADF)')
    p.add_argument('--elapsed',type=bounded(float,0,False),default=0,help='seconds since the last outbound packet (default 0)')
    p.add_argument('--lifetime',type=bounded(float,0,True),help='state lifetime in seconds; omit for no expiry')
    p=labs.add_parser('L10',help='tunnel overhead budget')
    p.add_argument('--path-mtu',type=bounded(int,0,True),default=1500,help='outer path MTU in bytes (default 1500)')
    labs.add_parser('L11',help='synthetic latency statistics')
    args=parser.parse_args()
    try:
        if args.lab=='L02':result=L02(args.window,args.rtt)
        elif args.lab=='L03':result=L03(args.route,args.lookup)
        elif args.lab=='L04':result=L04(args.ttl)
        elif args.lab=='L06':result=L06(args.window,not args.no_data_loss,not args.no_ack_loss)
        elif args.lab=='L07':result=L07(args.bc_cost)
        elif args.lab=='L08':result=L08(args.filter,args.elapsed,args.lifetime)
        elif args.lab=='L10':result=L10(args.path_mtu)
        else:result=L11()
    except ValueError as exc:parser.error(str(exc))
    except OverflowError:parser.error('input is outside the numeric range of this model')
    try:text=json.dumps(display(result),ensure_ascii=False,indent=2,allow_nan=False)
    except (ValueError,OverflowError):parser.error('result is outside the numeric range of this model; use smaller inputs')
    print(text)


if __name__=='__main__':main()