#!/usr/bin/env python3
# -*- coding: utf-8; py-indent-offset: 4 -*-
#
# Author:  Linuxfabrik GmbH, Zurich, Switzerland
# Contact: info (at) linuxfabrik (dot) ch
#          https://www.linuxfabrik.ch/
# License: The Unlicense, see LICENSE file.

# https://github.com/Linuxfabrik/monitoring-plugins/blob/main/CONTRIBUTING.md

"""See the check's README for more details."""

import argparse
import ipaddress
import socket
import struct
import sys
from random import randint
from uuid import getnode

import lib.args
import lib.base
from lib.globals import STATE_OK, STATE_UNKNOWN, STATE_WARN

__author__ = 'Linuxfabrik GmbH, Zurich/Switzerland'
__version__ = '2026040801'

DESCRIPTION = """Tests if a DHCP server can offer IPv4 addresses by emulating a DHCP client. Sends a
DHCPDISCOVER packet and verifies that the server responds with a valid DHCPOFFER. Only
performs the discovery step without requesting an actual lease (no DHCPREQUEST). Works
with both local and relayed DHCP servers, and can target a specific subnet.
Alerts if the server does not respond or the response is invalid.
Requires root or sudo."""

# DHCP relay must listen on all interfaces to receive broadcast DISCOVER
# packets from any client-facing subnet.
DEFAULT_BIND_ADDRESS = '0.0.0.0'  # nosec B104
DEFAULT_TIMEOUT = 7  # seconds


def parse_args():
    """Parse command line arguments using argparse."""
    parser = argparse.ArgumentParser(description=DESCRIPTION)

    parser.add_argument(
        '-V',
        '--version',
        action='version',
        version=f'%(prog)s: v{__version__} by {__author__}',
    )

    parser.add_argument(
        '--always-ok',
        help=lib.args.help('--always-ok'),
        dest='ALWAYS_OK',
        action='store_true',
        default=False,
    )

    parser.add_argument(
        '--bind-address',
        help='Local address to bind the socket to. Default: %(default)s',
        dest='BIND_ADDRESS',
        default=DEFAULT_BIND_ADDRESS,
    )

    parser.add_argument(
        '-H',
        '--hostname',
        help='DHCP server address, hostname or IP address. '
        'If omitted, the request is sent as broadcast. '
        'Default: %(default)s',
        dest='HOSTNAME',
        default=None,
    )

    parser.add_argument(
        '--mac',
        help='Network MAC address to use in the DHCP request. '
        'Does not have to be an existing MAC address. '
        'Use `--mac=random` for a random MAC address. '
        'If omitted, the local hardware address is used.',
        dest='MAC',
    )

    parser.add_argument(
        '--subnet-mask',
        help='Subnet mask for the DHCP request. '
        'Example: `255.255.255.248`. '
        'Default: %(default)s',
        dest='SUBNET_MASK',
        default=None,
    )

    parser.add_argument(
        '--subnet-selection',
        help='Override the DHCP server subnet selection for address allocation (RFC 3011). '
        'Example: `192.168.122.0`. '
        'Default: %(default)s',
        dest='SUBNET_SELECTION',
        default=None,
    )

    parser.add_argument(
        '--timeout',
        help=lib.args.help('--timeout') + ' Default: %(default)s (seconds)',
        dest='TIMEOUT',
        type=int,
        default=DEFAULT_TIMEOUT,
    )

    args, _ = parser.parse_known_args()
    return args


def get_transaction_id():
    """Get a random 4-byte transaction ID."""
    xid = b''
    for _ in range(4):
        xid += struct.pack('!B', randint(0, 255))
    return xid


def randmac():
    """Return random MAC address (6 bytes)."""
    macb = b''
    for _ in range(6):
        macb += struct.pack('!B', randint(0, 255))
    return macb


def normalize_mac(mac):
    """Normalize a given MAC address."""
    mac = str(mac).replace(':', '')
    if mac.startswith('0x'):
        # remove "hex()" prefix
        mac = mac[2:]
    while len(mac) < 12:
        mac = '0' + mac
    return mac


def mac2bytes(mac):
    """Convert a normalized hardware address (str) to an 48-bit positive integer in byte format."""
    macb = b''
    for i in range(0, 12, 2):
        macb += struct.pack('!B', int(mac[i : i + 2], 16))
    return macb


def create_dhcp_discover_msg(macb, xid, subnet_mask, subnet_selection):
    """Create a DHCP discover message.

    When a client needs an IP address for the first time, it sends a DHCPDISCOVER message
    (with its MAC address) as a network broadcast to the available DHCP servers (there may
    be several in the same subnet). This broadcast has 0.0.0.0 as the sender IP address and
    255.255.255.255 as the destination address, since the sender does not yet have an IP
    address and is sending its request "to all". The source UDP port is 68 and the destination
    UDP port is 67. The DHCP servers respond with DHCPOFFER and make suggestions for an IP
    address. This is done either with a broadcast to the address 255.255.255.255 with UDP
    source port 67 and UDP destination port 68, or with a unicast to the suggested IP address
    and the MAC address of the client, depending on whether the client has set the broadcast
    bit in the DHCPDISCOVER message.

       FIELD      OCTETS       DESCRIPTION
       -----      ------       -----------

       op            1  Message op code / message type.
                        1 = BOOTREQUEST, 2 = BOOTREPLY
       htype         1  Hardware address type, see ARP section in "Assigned
                        Numbers" RFC; e.g., '1' = 10mb ethernet.
       hlen          1  Hardware address length (e.g.  '6' for 10mb
                        ethernet).
       hops          1  Client sets to zero, optionally used by relay agents
                        when booting via a relay agent.
       xid           4  Transaction ID, a random number chosen by the
                        client, used by the client and server to associate
                        messages and responses between a client and a
                        server.
       secs          2  Filled in by client, seconds elapsed since client
                        began address acquisition or renewal process.
       flags         2  Flags (see figure 2).
       ciaddr        4  Client IP address; only filled in if client is in
                        BOUND, RENEW or REBINDING state and can respond
                        to ARP requests.
       yiaddr        4  'your' (client) IP address.
       siaddr        4  IP address of next server to use in bootstrap;
                        returned in DHCPOFFER, DHCPACK by server.
       giaddr        4  Relay agent IP address, used in booting via a
                        relay agent.
       chaddr       16  Client hardware address.
       sname        64  Optional server host name, null terminated string.
       file        128  Boot file name, null terminated string; "generic"
                        name or null in DHCPDISCOVER, fully qualified
                        directory-path name in DHCPOFFER.
       options     var  Optional parameters field.  See the options
                        documents for a list of defined options.
    """

    # ! = network byte order which is always big-endian
    # B = 1 byte unsigned char, H = 2 bytes unsigned short, I = 4 bytes unsigned int           Bytes
    dhcp_discover = struct.pack(
        '!B', 1
    )  # Op Message type: 1 = Boot Request                    43
    dhcp_discover += struct.pack(
        '!B', 1
    )  # HTYPE: Ethernet                                      44
    dhcp_discover += struct.pack(
        '!B', 6
    )  # HLEN                                                 45
    dhcp_discover += struct.pack(
        '!B', 0
    )  # HOPS                                                 46
    dhcp_discover += (
        xid  # Transaction ID                                                    47..50
    )
    dhcp_discover += struct.pack(
        '!H', 0
    )  # SECS                                              51 52
    dhcp_discover += struct.pack(
        '!H', 0x8000
    )  # FLAGS: 8000h = Broadcast                     53 54
    dhcp_discover += struct.pack(
        '!I', 0
    )  # CIADDR: 0.0.0.0                                  55..58
    dhcp_discover += struct.pack(
        '!I', 0
    )  # YIADDR: 0.0.0.0                                  59..62
    dhcp_discover += struct.pack(
        '!I', 0
    )  # SIADDR: 0.0.0.0                                  63..66
    dhcp_discover += struct.pack(
        '!I', 0
    )  # GIADDR: 0.0.0.0                                  67..70
    dhcp_discover += (
        macb  #                                                                   71..76
    )
    dhcp_discover += (
        struct.pack('!B', 0) * 10
    )  # Client hardware address padding              77..86
    dhcp_discover += (
        struct.pack('!B', 0) * 64
    )  # SNAME: Server host name not given           87..150
    dhcp_discover += (
        struct.pack('!B', 0) * 128
    )  # FILE: Boot file name not given            151..278

    # dhcp options: https://datatracker.ietf.org/doc/html/rfc2132
    # everything that follows the "Magic Cookie 0x63825363" is to be interpreted as DHCP options
    # https://datatracker.ietf.org/doc/html/rfc951
    dhcp_discover += struct.pack(
        '!I', 0x63825363
    )  #                                        279..282

    dhcp_discover += struct.pack('!B', 0x35) + struct.pack(
        '!H', 0x0101
    )  # code 53 (dec): DHCP Message Type = DHCP discover
    dhcp_discover += (
        struct.pack('!H', 0x3D06) + macb
    )  # code 61 (dec): Client-identifier, 6 bytes long
    if subnet_mask:
        # https://datatracker.ietf.org/doc/html/rfc2132#section-3.3
        dhcp_discover += struct.pack(
            '!H', 0x0104
        )  # code 01 (dec): Subnet Mask, 4 bytes
        try:
            dhcp_discover += int(ipaddress.IPv4Address(subnet_mask)).to_bytes(
                4, byteorder='big'
            )
        except ipaddress.AddressValueError:
            lib.base.oao(f'Invaild subnet mask "{subnet_mask}"', STATE_UNKNOWN)
    if subnet_selection:
        # https://www.rfc-editor.org/rfc/rfc3011.html
        dhcp_discover += struct.pack(
            '!H', 0x7604
        )  # code 118 (dec): Subnet Selection, 4 bytes
        try:
            dhcp_discover += int(ipaddress.IPv4Address(subnet_selection)).to_bytes(
                4, byteorder='big'
            )
        except ipaddress.AddressValueError:
            lib.base.oao(
                f'Invaild subnet selection "{subnet_selection}"', STATE_UNKNOWN
            )
    dhcp_discover += struct.pack('!I', 0xFF)  # code 255 (dec): End option

    return dhcp_discover


def decode_dhcp_options(packet):
    """Return some info from DHCP options, right after the Magic cookie."""
    # find magic cookie 99.130.83.99 (or hexadecimal number 63.82.53.63)
    mc_pos = packet.find(b'\x63\x82\x53\x63') + 4

    # https://datatracker.ietf.org/doc/html/rfc2132#section-3.3
    pos = packet.find(b'\x01\x04', mc_pos) + 2  # option 1, len 4
    subnet_mask = '.'.join(str(byte) for byte in packet[pos : pos + 4])

    # https://datatracker.ietf.org/doc/html/rfc2132#section-5.3
    pos = packet.find(b'\x1c\x04', mc_pos) + 2  # option 28, len 4
    bc_addr = '.'.join(str(byte) for byte in packet[pos : pos + 4])

    return subnet_mask, bc_addr


def main():
    """The main function. This is where the magic happens."""

    # parse the command line
    try:
        args = parse_args()
    except SystemExit:
        sys.exit(STATE_UNKNOWN)

    dhcps = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)  # Internet, UDP
    dhcps.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    try:
        dhcps.bind((args.BIND_ADDRESS, 68))  # we want to send from port 68
    except PermissionError as e:
        dhcps.close()
        lib.base.oao(f'Permission error "{e.strerror}"', STATE_UNKNOWN)
    except OSError as e:
        dhcps.close()
        lib.base.oao(f'OS error "{e.strerror}"', STATE_UNKNOWN)

    # handle MAC address
    if args.MAC:
        if args.MAC.lower().startswith('r'):
            mac = 'rand'
            macb = randmac()
        else:
            mac = normalize_mac(args.MAC)
            macb = mac2bytes(mac)
    else:
        mac = normalize_mac(hex(getnode()))
        macb = mac2bytes(mac)

    # fetch data
    xid = get_transaction_id()
    dhcp_discover = create_dhcp_discover_msg(
        macb, xid, args.SUBNET_MASK, args.SUBNET_SELECTION
    )
    dhcps.settimeout(args.TIMEOUT)
    if args.HOSTNAME:
        dhcps.sendto(dhcp_discover, (args.HOSTNAME, 67))
    else:
        dhcps.sendto(dhcp_discover, ('<broadcast>', 67))
    try:
        while True:
            dhcp_offer = dhcps.recv(1024)
            if dhcp_offer[4:8] == xid:
                # extract addresses
                yiaddr = '.'.join(str(byte) for byte in dhcp_offer[16:20])
                siaddr = '.'.join(str(byte) for byte in dhcp_offer[20:24])
                # and some options
                subnet_mask, bc_addr = decode_dhcp_options(dhcp_offer)
                break  # jump out of the loop, otherwise we'll get a "timed out" message
        dhcps.close()
    except socket.timeout:
        lib.base.oao(
            'Socket timeout. Possibly DHCP pool is exhausted, does not exist, or similar.',
            STATE_WARN,
        )

    # analyze data
    state = STATE_OK
    # String comparison against the DHCP wire-format "not assigned" sentinel,
    # not a bind address.
    if yiaddr == '0.0.0.0':  # nosec B104
        state = STATE_WARN

    # build the message
    msg = (
        f'DHCPOFFER: IP={yiaddr}/{subnet_mask}'
        f'{lib.base.state2str(state, prefix=" ")}'
        f' Server ID={siaddr}'
        f' Broadcast Addr={bc_addr}\n'
    )
    msg += (
        f'DHCPDISCOVER: MAC={mac}'
        f' Host={args.HOSTNAME if args.HOSTNAME else "Broadcast"}'
        f' Network={args.SUBNET_SELECTION}'
        f'/{args.SUBNET_MASK}'
    )

    # over and out
    lib.base.oao(msg, state, always_ok=args.ALWAYS_OK)


if __name__ == '__main__':
    try:
        main()
    except Exception:
        lib.base.cu()
