#!/usr/bin/env python3
"""
轻量 DHCP Server — 用于 CH395F Phase 6 测试
用法：以管理员身份运行
    python dhcp_server.py --iface 192.168.1.2
"""
import socket
import struct
import argparse
import time

DHCP_SERVER_PORT = 67
DHCP_CLIENT_PORT = 68

# DHCP 消息类型
DHCP_DISCOVER = 1
DHCP_OFFER = 2
DHCP_REQUEST = 3
DHCP_ACK = 5

# DHCP 选项
OPT_MESSAGE_TYPE = 53
OPT_SERVER_ID = 54
OPT_REQUESTED_IP = 50
OPT_SUBNET_MASK = 1
OPT_ROUTER = 3
OPT_DNS = 6
OPT_LEASE_TIME = 51
OPT_END = 255

SERVER_IP = "192.168.1.2"
OFFER_IP = "192.168.1.100"
SUBNET_MASK = "255.255.255.0"
ROUTER = "192.168.1.1"
LEASE_TIME = 3600  # 1 小时

def ip_to_bytes(ip):
    return bytes(map(int, ip.split('.')))

def build_dhcp_offer(discover_pkt, server_ip, offer_ip):
    """构建 DHCP OFFER 包"""
    transaction_id = discover_pkt[232:236]  # Bootp header 前 4 字节
    client_mac = discover_pkt[28:34]

    # BOOTP 头
    op = b'\x02'           # Boot reply
    htype = b'\x01'        # Ethernet
    hlen = b'\x06'         # MAC len
    hops = b'\x00'
    xid = transaction_id
    secs = b'\x00\x00'
    flags = b'\x00\x00'
    ciaddr = b'\x00\x00\x00\x00'
    yiaddr = ip_to_bytes(offer_ip)   # Your IP
    siaddr = ip_to_bytes(server_ip)  # Server IP
    giaddr = b'\x00\x00\x00\x00'
    chaddr = client_mac + b'\x00' * 10  # Client MAC (padded to 16)
    sname = b'\x00' * 64
    file = b'\x00' * 128
    magic = b'\x63\x82\x53\x63'  # DHCP magic cookie

    # DHCP 选项
    options = b''
    options += struct.pack('BBB', OPT_MESSAGE_TYPE, 1, DHCP_OFFER)
    options += struct.pack('BB', OPT_SERVER_ID, 4) + ip_to_bytes(server_ip)
    options += struct.pack('BB', OPT_REQUESTED_IP, 4) + ip_to_bytes(offer_ip)
    options += struct.pack('BB', OPT_SUBNET_MASK, 4) + ip_to_bytes(SUBNET_MASK)
    options += struct.pack('BB', OPT_ROUTER, 4) + ip_to_bytes(ROUTER)
    options += struct.pack('BB', OPT_DNS, 4) + ip_to_bytes(ROUTER)
    options += bytes([OPT_LEASE_TIME, 4]) + struct.pack('!I', LEASE_TIME)
    options += bytes([OPT_END])

    pkt = op + htype + hlen + hops + xid + secs + flags + ciaddr + yiaddr + siaddr + giaddr + chaddr + sname + file + magic + options
    return pkt

def build_dhcp_ack(request_pkt, server_ip, offer_ip):
    """构建 DHCP ACK 包"""
    transaction_id = request_pkt[232:236]
    client_mac = request_pkt[28:34]

    op = b'\x02'
    htype = b'\x01'
    hlen = b'\x06'
    hops = b'\x00'
    xid = transaction_id
    secs = b'\x00\x00'
    flags = b'\x00\x00'
    ciaddr = b'\x00\x00\x00\x00'
    yiaddr = ip_to_bytes(offer_ip)
    siaddr = ip_to_bytes(server_ip)
    giaddr = b'\x00\x00\x00\x00'
    chaddr = client_mac + b'\x00' * 10
    sname = b'\x00' * 64
    file = b'\x00' * 128
    magic = b'\x63\x82\x53\x63'

    options = b''
    options += struct.pack('BBB', OPT_MESSAGE_TYPE, 1, DHCP_ACK)
    options += struct.pack('BB', OPT_SERVER_ID, 4) + ip_to_bytes(server_ip)
    options += struct.pack('BB', OPT_SUBNET_MASK, 4) + ip_to_bytes(SUBNET_MASK)
    options += struct.pack('BB', OPT_ROUTER, 4) + ip_to_bytes(ROUTER)
    options += struct.pack('BB', OPT_DNS, 4) + ip_to_bytes(ROUTER)
    options += bytes([OPT_LEASE_TIME, 4]) + struct.pack('!I', LEASE_TIME)
    options += bytes([OPT_END])

    return op + htype + hlen + hops + xid + secs + flags + ciaddr + yiaddr + siaddr + giaddr + chaddr + sname + file + magic + options

def main():
    parser = argparse.ArgumentParser(description="轻量 DHCP Server for CH395F 测试")
    parser.add_argument("--iface", default=SERVER_IP, help="网卡 IP 地址")
    args = parser.parse_args()

    offer_ip = OFFER_IP
    server_ip = args.iface

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    sock.bind(('0.0.0.0', DHCP_SERVER_PORT))
    sock.settimeout(60)

    print(f"DHCP Server 启动, 网卡 IP={server_ip}, 分配 IP={offer_ip}")
    print("等待 CH395F DHCP DISCOVER...")

    try:
        # 第一阶段：DISCOVER → OFFER
        data, addr = sock.recvfrom(4096)
        msg_type = data[232 + 4 + 1] if len(data) > 250 else 0
        client_mac = ':'.join(f'{b:02X}' for b in data[28:34])
        print(f"收到 DHCP DISCOVER from {client_mac}")

        offer = build_dhcp_offer(data, server_ip, offer_ip)
        sock.sendto(offer, ('255.255.255.255', DHCP_CLIENT_PORT))
        print("发送 DHCP OFFER")

        # 第二阶段：REQUEST → ACK
        sock.settimeout(10)
        try:
            data, addr = sock.recvfrom(4096)
            client_mac = ':'.join(f'{b:02X}' for b in data[28:34])
            print(f"收到 DHCP REQUEST from {client_mac}")

            ack = build_dhcp_ack(data, server_ip, offer_ip)
            sock.sendto(ack, ('255.255.255.255', DHCP_CLIENT_PORT))
            print(f"发送 DHCP ACK → {offer_ip}")
            print("DHCP 分配完成！")
        except socket.timeout:
            print("等待 DHCP REQUEST 超时")

    except socket.timeout:
        print("等待 DHCP DISCOVER 超时（60s）")
    finally:
        sock.close()

if __name__ == '__main__':
    main()