pxe-server/tests/pxe-nic-check/frames.py

"""
frames.py - synthesize the Ethernet frames pxe-nic-check looks for.

Shared by the unit tests (fed straight to the parsers) and the namespace
lab (sent on a veth by lab_peer.py). Plain struct packing, no scapy.
"""

import os
import struct
import socket
import importlib.machinery
import importlib.util

HERE = os.path.dirname(os.path.abspath(__file__))
SCRIPT = os.path.join(HERE, "..", "..", "playbook", "pxe-server-helpers",
                      "pxe-nic-check")


def load_script():
    """Import the extensionless script as a module (no __pycache__ next
    to it: that directory is copied to the server)."""
    import sys
    sys.dont_write_bytecode = True
    loader = importlib.machinery.SourceFileLoader("pxe_nic_check", SCRIPT)
    spec = importlib.util.spec_from_loader("pxe_nic_check", loader)
    mod = importlib.util.module_from_spec(spec)
    loader.exec_module(mod)
    return mod


def mac(s):
    return bytes.fromhex(s.replace(":", ""))


def eth(dst, src, etype, payload):
    return mac(dst) + mac(src) + struct.pack("!H", etype) + payload


def _tlv(t, value):
    return struct.pack("!H", (t << 9) | len(value)) + value


def lldp(src, system="WJ-IDF2", port="Gi1/0/14", vlan=None,
         chassis_mac="00:11:22:33:44:55"):
    body = _tlv(1, b"\x04" + mac(chassis_mac))          # chassis: MAC
    body += _tlv(2, b"\x05" + port.encode())             # port: ifName
    body += _tlv(3, struct.pack("!H", 120))              # TTL
    if system is not None:
        body += _tlv(5, system.encode())
    if vlan is not None:
        body += _tlv(127, b"\x00\x80\xc2\x01" + struct.pack("!H", vlan))
    body += _tlv(0, b"")
    return eth("01:80:c2:00:00:0e", src, 0x88CC, body)


def cdp(src, device="WJ-CORE1.wjs.geaerospace.net", port="GigabitEthernet1/0/3",
        vlan=None):
    def t(code, val):
        return struct.pack("!HH", code, 4 + len(val)) + val
    tlvs = t(1, device.encode()) + t(3, port.encode())
    if vlan is not None:
        tlvs += t(0x0A, struct.pack("!H", vlan))
    body = bytes([2, 180]) + b"\x00\x00" + tlvs
    llc = b"\xaa\xaa\x03" + b"\x00\x00\x0c" + b"\x20\x00" + body
    return mac("01:00:0c:cc:cc:cc") + mac(src) + struct.pack("!H", len(llc)) \
        + llc


def eapol(src):
    # EAPOL v2, type 0 (EAP packet), EAP Request/Identity
    eap = bytes([1, 1]) + struct.pack("!H", 5) + bytes([1])
    body = bytes([2, 0]) + struct.pack("!H", len(eap)) + eap
    return eth("01:80:c2:00:00:03", src, 0x888E, body)


def stp(src):
    bpdu = b"\x00\x00\x00\x00\x00" + b"\x80\x00" + mac(src) + b"\x00" * 23
    llc = b"\x42\x42\x03" + bpdu
    return mac("01:80:c2:00:00:00") + mac(src) + \
        struct.pack("!H", len(llc)) + llc


def arp(src, sender_ip="172.16.9.50", target_ip="172.16.9.1"):
    body = struct.pack("!HHBBH", 1, 0x0800, 6, 4, 1) + mac(src) + \
        socket.inet_aton(sender_ip) + b"\x00" * 6 + socket.inet_aton(target_ip)
    return eth("ff:ff:ff:ff:ff:ff", src, 0x0806, body)


def _dns_encode(names):
    out = b""
    for n in names:
        for label in n.split("."):
            out += bytes([len(label)]) + label.encode()
        out += b"\x00"
    return out


def _ip_checksum(hdr):
    s = 0
    for i in range(0, len(hdr), 2):
        s += (hdr[i] << 8) + hdr[i + 1]
    while s >> 16:
        s = (s & 0xFFFF) + (s >> 16)
    return (~s) & 0xFFFF


def dhcp_reply(src, server_ip, xid, chaddr, yiaddr="10.99.0.50",
               domain=None, search=None, msgtype=2, server_id=None,
               search_raw=None):
    """DHCPOFFER (msgtype 2) or ACK (5), broadcast."""
    sid = server_ip if server_id is None else server_id
    bootp = struct.pack("!BBBBIHH", 2, 1, 6, 0, xid, 0, 0x8000)
    bootp += b"\x00" * 4 + socket.inet_aton(yiaddr) + \
        socket.inet_aton(server_ip) + b"\x00" * 4
    bootp += mac(chaddr) + b"\x00" * 10 + b"\x00" * 64 + b"\x00" * 128
    opts = b"\x63\x82\x53\x63" + bytes([53, 1, msgtype])
    if sid:
        opts += bytes([54, 4]) + socket.inet_aton(sid)
    opts += bytes([1, 4]) + socket.inet_aton("255.255.255.0")
    opts += bytes([51, 4]) + struct.pack("!I", 3600)
    if domain:
        opts += bytes([15, len(domain)]) + domain.encode()
    if search or search_raw:
        raw = search_raw if search_raw is not None else _dns_encode(search)
        opts += bytes([119, len(raw)]) + raw
    opts += b"\xff"
    bootp += opts
    if len(bootp) < 300:
        bootp += b"\x00" * (300 - len(bootp))
    udp = struct.pack("!HHHH", 67, 68, 8 + len(bootp), 0) + bootp
    ip = struct.pack("!BBHHHBBH4s4s", 0x45, 0, 20 + len(udp), 1, 0, 64, 17,
                     0, socket.inet_aton(server_ip),
                     socket.inet_aton("255.255.255.255"))
    ip = ip[:10] + struct.pack("!H", _ip_checksum(ip)) + ip[12:]
    return eth("ff:ff:ff:ff:ff:ff", src, 0x0800, ip + udp)


def parse_discover(frame):
    """(xid, chaddr) if frame is a DHCPDISCOVER, else None."""
    if len(frame) < 14 + 20 + 8 + 240 or frame[12:14] != b"\x08\x00":
        return None
    ip = frame[14:]
    ihl = (ip[0] & 0x0F) * 4
    if ip[9] != 17:
        return None
    sport, dport = struct.unpack("!HH", ip[ihl:ihl + 4])
    if (sport, dport) != (68, 67):
        return None
    b = ip[ihl + 8:]
    if b[0] != 1 or b[236:240] != b"\x63\x82\x53\x63":
        return None
    opts = b[240:]
    i = 0
    mtype = None
    while i < len(opts) and opts[i] != 255:
        if opts[i] == 0:
            i += 1
            continue
        if opts[i] == 53:
            mtype = opts[i + 2]
        i += 2 + opts[i + 1]
    if mtype != 1:
        return None
    xid = struct.unpack("!I", b[4:8])[0]
    chaddr = ":".join("%02x" % x for x in b[28:34])
    return xid, chaddr