#!/usr/bin/env python3
# dhcpsearch -- generate DHCP option 119 (domain search list)
import argparse
import io
import struct

class DnsPacketIO(io.BytesIO):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._suffixes = {}

    def write_u8(self, value):
        return self.write(struct.pack("B", value))

    def write_u16_be(self, value):
        return self.write(struct.pack(">H", value))

    def write_domain(self, domain):
        labels = domain.encode().lower().strip(b".").split(b".") + [b""]
        ret = 0
        for i, label in enumerate(labels):
            suffix = b".".join(labels[i:]).lower()
            if len(suffix) > 0:
                # Note: Suffix lookup for compression must be case-insensitive.
                if suffix in self._suffixes:
                    offset = self._suffixes[suffix]
                    ret += self.write_u16_be(0xC000 | offset)
                    break
                else:
                    self._suffixes[suffix] = self.tell()
            ret += self.write_u8(len(label))
            ret += self.write(label)
        return ret

parser = argparse.ArgumentParser()
parser.add_argument("domain", nargs="+",
                    help="domain suffix")
args = parser.parse_args()

print("Domains:")
pkt = DnsPacketIO()
for domain in args.domain:
    print(f"\t{domain}")
    pkt.write_domain(domain)

buf = pkt.getvalue()

print()
print(f"Option code: 119")
print(f"Option data: {buf!r}")
print()
print(f"Mikrotik RouterOS:")
print(f"/ip dhcp-server option")
print(f"add name=domain-search code=119 value=0x{buf.hex().upper()}")
