package main

import (
	"bytes"
	"encoding/binary"
	"net"
	"unsafe"

	"go.uber.org/zap"
)

type DhcpOp uint8

const (
	BOOTREQUEST DhcpOp = iota + 1
	BOOTREPLY
)

type DhcpMessage struct {
	Op     DhcpOp
	Htype  uint8
	Hlen   uint8
	Hops   uint8
	Xid    uint32
	Secs   uint16
	Flags  uint16
	Ciaddr uint32
	Yiaddr uint32
	Siaddr uint32
	Giaddr uint32
	Chaddr [16]byte
	Sname  [64]byte
	File   [128]byte
	Magic  [4]byte
}

func (dhcp DhcpMessage) Debug(s *zap.SugaredLogger) {
	s.Debugf("op=%x, htype=%x, hlen=%x, hops=%x", dhcp.Op, dhcp.Htype, dhcp.Hlen, dhcp.Hops)
	s.Debugf("xid=%x", dhcp.Xid)
	s.Debugf("secs=%d, flags=%x", dhcp.Secs, dhcp.Flags)

	s.Debugf("ciaddr=%x", dhcp.Ciaddr)
	s.Debugf("siaddr=%x", dhcp.Siaddr)
	s.Debugf("giaddr=%x", dhcp.Giaddr)
	s.Debugf("chaddr=%x", dhcp.Chaddr)

	s.Debugf("sname=%s", string(dhcp.Sname[:]))
	s.Debugf("file=%s", string(dhcp.File[:]))
	s.Debugf("magic=%x", dhcp.Magic)
}

func parseOptions(data []byte) {
	s := zap.S()
	i := 0
	for i < len(data) {
		code := data[i]
		if code == 255 {
			s.Debug("Found END option at offset ", i)
			break
		}

		if code == 0 {
			s.Debug("Found padding option at offset ", i)
			i += 1
			continue
		}

		size := int(data[i+1])
		payload := data[i+2 : i+2+size]

		s.Debugf("code=%d, size=%d, payload=%x", code, size, payload)
		i += size + 2
	}
}

func prepareOffer(request DhcpMessage) DhcpMessage {
	var sname [64]byte
	var file [128]byte
	copy(sname[:], "go-dhcp-server")

	dhcp := DhcpMessage{
		Op:     BOOTREPLY,
		Htype:  1, // Ethernet
		Hlen:   6, // Ethernet address length
		Hops:   0,
		Secs:   0,
		Flags:  request.Flags,
		Xid:    request.Xid,
		Siaddr: binary.BigEndian.Uint32(net.IPv4(10, 0, 0, 1).To4()),
		Ciaddr: 0,
		Yiaddr: binary.BigEndian.Uint32(net.IPv4(10, 0, 0, 2).To4()),
		Giaddr: request.Giaddr,
		Chaddr: request.Chaddr,
		Magic:  [4]byte{0x63, 0x82, 0x53, 0x63},
		Sname:  sname,
		File:   file,
	}
	return dhcp
}

func handleMsg(data []byte, remote *net.UDPAddr) {
	s := zap.S()

	s.Debugf("Connection from client %v", remote.IP)

	var dhcp DhcpMessage
	reader := bytes.NewReader(data)
	binary.Read(reader, binary.BigEndian, &dhcp)
	dhcp.Debug(s)

	if dhcp.Magic != [4]byte{0x63, 0x82, 0x53, 0x63} {
		panic("Invalid DHCP magic field")
	}

	optDataOffset := int(unsafe.Sizeof(dhcp))
	optData := data[optDataOffset:]
	s.Debug("Raw options data:", optData)

	parseOptions(optData)

	offer := prepareOffer(dhcp)
	localAddr, _ := net.ResolveUDPAddr("udp", "172.17.0.1:68")
	clientAddr, _ := net.ResolveUDPAddr("udp", "255.255.255.255:68")
	conn, err := net.DialUDP("udp", localAddr, clientAddr)
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	s.Info("Sending DHCPOFFER...")
	buf := make([]byte, 0)
	w := bytes.NewBuffer(buf)
	binary.Write(w, binary.BigEndian, offer)
	w.Write([]byte{53, 1, 2})
	w.WriteByte(255)
	msg := w.Bytes()
	s.Debug("Msg", msg)
	conn.Write(msg)
}

func main() {
	// Set up logging
	logger, _ := zap.NewDevelopment()
	defer logger.Sync() // flushes buffer, if any
	zap.ReplaceGlobals(logger)

	addr, _ := net.ResolveUDPAddr("udp4", ":67")
	conn, err := net.ListenUDP("udp4", addr)
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	zap.L().Info("Listening for incoming connections")

	for {
		buf := make([]byte, 1024)
		rlen, remote, err := conn.ReadFromUDP(buf[:])
		if err != nil {
			panic(err)
		}

		// Do stuff with the read bytes
		handleMsg(buf[0:rlen], remote)
	}
}