"""One household device's DNS traffic, stdlib only.

usage: traffic.py PROFILE bulk          send the profile's lookup budget fast, then exit
       traffic.py PROFILE live MEAN_S   one lookup every ~MEAN_S seconds, forever
"""
import os
import random
import socket
import struct
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from profiles import PROFILES  # noqa: E402

SERVER = (os.environ.get("NXDNS", "192.168.50.2"), 53)


def build_query(qid, name, qtype):
    header = struct.pack("!HHHHHH", qid, 0x0100, 1, 0, 0, 1)
    qname = b"".join(bytes([len(label)]) + label.encode() for label in name.split(".")) + b"\0"
    opt = b"\0" + struct.pack("!HHIH", 41, 1232, 0, 0)
    return header + qname + struct.pack("!HH", qtype, 1) + opt


def pick_types(type_sets):
    sets, weights = zip(*type_sets)
    return random.choices(sets, weights)[0]


def main():
    name, mode = sys.argv[1], sys.argv[2]
    profile = PROFILES[name]
    domains, weights = zip(*profile["domains"])
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(6)
    sent = answered = failed = 0
    lookups = profile["lookups"] if mode == "bulk" else None
    mean = float(sys.argv[3]) if mode == "live" else 0.0
    done = 0
    while lookups is None or done < lookups:
        domain = random.choices(domains, weights)[0]
        for qtype in pick_types(profile["types"]):
            qid = random.randrange(65536)
            sock.sendto(build_query(qid, domain, qtype), SERVER)
            sent += 1
            try:
                while True:
                    data, _ = sock.recvfrom(4096)
                    if struct.unpack("!H", data[:2])[0] == qid:
                        answered += 1
                        break
            except socket.timeout:
                failed += 1
                print(f"{name}: no answer for {domain} type {qtype}", flush=True)
        done += 1
        if mode == "bulk":
            time.sleep(random.uniform(0.05, 0.25))
        else:
            time.sleep(random.expovariate(1.0 / mean))
        if done % 500 == 0:
            print(f"{name}: {done} lookups, {sent} sent, {answered} answered, {failed} unanswered", flush=True)
    print(f"{name}: finished {done} lookups, {sent} sent, {answered} answered, {failed} unanswered", flush=True)


if __name__ == "__main__":
    main()