feat: Turn into Cobra application, add CLI flag for listen interface
This commit is contained in:
Adrian Rumpold
330
cmd/root.go
Normal file
330
cmd/root.go
Normal file
@@ -0,0 +1,330 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"unsafe"
|
||||
|
||||
"github.com/adrianokf/go-dhcp/pkg/leases"
|
||||
"github.com/adrianokf/go-dhcp/pkg/messages"
|
||||
"github.com/adrianokf/go-dhcp/pkg/types"
|
||||
"github.com/spf13/cobra"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
var listenInterface string
|
||||
|
||||
var manager = leases.NewLeaseManager()
|
||||
|
||||
var magic = [4]byte{0x63, 0x82, 0x53, 0x63}
|
||||
|
||||
func parseOptions(data []byte) messages.Options {
|
||||
s := zap.S()
|
||||
i := 0
|
||||
|
||||
options := make(messages.Options)
|
||||
|
||||
out:
|
||||
for i < len(data) {
|
||||
code := messages.OptionCode(data[i])
|
||||
switch code {
|
||||
case messages.OptionEnd:
|
||||
s.Debug("Found END option at offset ", i)
|
||||
break out
|
||||
|
||||
case messages.OptionPad:
|
||||
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)
|
||||
options[code] = messages.Option{Code: code, Data: data[i+1 : i+2+size]}
|
||||
i += size + 2
|
||||
}
|
||||
|
||||
s.Debugf("Parsed options: ", options)
|
||||
return options
|
||||
}
|
||||
|
||||
func prepareOffer(request messages.DhcpMessage, lease leases.Lease) messages.DhcpMessage {
|
||||
var sname [64]byte
|
||||
var file [128]byte
|
||||
var siaddr types.Ipv4Addr
|
||||
copy(sname[:], "go-dhcp-server")
|
||||
copy(siaddr[:], net.IPv4(10, 0, 0, 1).To4())
|
||||
|
||||
dhcp := messages.DhcpMessage{
|
||||
Op: messages.BOOTREPLY,
|
||||
Htype: 1, // Ethernet
|
||||
Hlen: 6, // Ethernet address length
|
||||
Hops: 0,
|
||||
Secs: 0,
|
||||
Flags: request.Flags,
|
||||
Xid: lease.TransactionId,
|
||||
Siaddr: siaddr,
|
||||
Ciaddr: [4]byte{0, 0, 0, 0},
|
||||
Yiaddr: lease.Address,
|
||||
Giaddr: request.Giaddr,
|
||||
Chaddr: request.Chaddr,
|
||||
Magic: magic,
|
||||
Sname: sname,
|
||||
File: file,
|
||||
}
|
||||
return dhcp
|
||||
}
|
||||
|
||||
func prepareAck(request messages.DhcpMessage, lease leases.Lease) messages.DhcpMessage {
|
||||
var sname [64]byte
|
||||
var file [128]byte
|
||||
copy(sname[:], "go-dhcp-server")
|
||||
|
||||
dhcp := messages.DhcpMessage{
|
||||
Op: messages.BOOTREPLY,
|
||||
Htype: 1, // Ethernet
|
||||
Hlen: 6, // Ethernet address length
|
||||
Hops: 0,
|
||||
Secs: 0,
|
||||
Flags: 0,
|
||||
Xid: request.Xid,
|
||||
Siaddr: lease.Address,
|
||||
Ciaddr: request.Ciaddr,
|
||||
Yiaddr: lease.Address,
|
||||
Giaddr: request.Giaddr,
|
||||
Chaddr: request.Chaddr,
|
||||
Magic: magic,
|
||||
Sname: sname,
|
||||
File: file,
|
||||
}
|
||||
return dhcp
|
||||
}
|
||||
|
||||
// sendMessage transmits a DHCP message with options via a UDP connection
|
||||
// The end option (code 255) is automatically appended and does not need to
|
||||
// be passed explicitly.
|
||||
func sendMessage(conn *net.UDPConn, message messages.DhcpMessage, options []messages.Option) error {
|
||||
buf := make([]byte, 0)
|
||||
w := bytes.NewBuffer(buf)
|
||||
err := binary.Write(w, binary.BigEndian, message)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, v := range options {
|
||||
err = w.WriteByte(byte(v.Code))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = w.Write(v.Data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Automatically add END option, so the caller doesn't
|
||||
// need to specificy it for every invocation.
|
||||
err = w.WriteByte(byte(messages.OptionEnd))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
msg := w.Bytes()
|
||||
zap.S().Debug("Msg", msg)
|
||||
_, err = conn.Write(msg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func handleOffer(dhcp messages.DhcpMessage, remote *net.UDPAddr) error {
|
||||
s := zap.S()
|
||||
|
||||
lease, err := manager.Request(dhcp.Xid, dhcp.Chaddr)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
offer := prepareOffer(dhcp, *lease)
|
||||
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...")
|
||||
options := []messages.Option{
|
||||
{
|
||||
Code: messages.OptionDHCPMessageType,
|
||||
Data: []byte{1, byte(messages.MessageTypeOffer)},
|
||||
},
|
||||
}
|
||||
sendMessage(conn, offer, options)
|
||||
return nil
|
||||
}
|
||||
|
||||
func handleAck(dhcp messages.DhcpMessage, remote *net.UDPAddr) error {
|
||||
s := zap.S()
|
||||
|
||||
lease, err := manager.Lookup(dhcp.Xid)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
ack := prepareAck(dhcp, *lease)
|
||||
|
||||
options := []messages.Option{
|
||||
{
|
||||
Code: messages.OptionDHCPMessageType,
|
||||
Data: []byte{1, byte(messages.MessageTypeAck)},
|
||||
},
|
||||
{
|
||||
Code: messages.OptionIPAddressLeaseTime,
|
||||
Data: append([]byte{4}, u32tob(3600)...),
|
||||
},
|
||||
}
|
||||
s.Debug("Options: ", options)
|
||||
|
||||
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 DHCPACK")
|
||||
sendMessage(conn, ack, options)
|
||||
|
||||
lease, err = manager.Request(dhcp.Xid, dhcp.Chaddr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.Debug("Found lease", lease)
|
||||
return nil
|
||||
}
|
||||
|
||||
func handleMsg(data []byte, remote *net.UDPAddr) {
|
||||
s := zap.S()
|
||||
|
||||
s.Debugf("Connection from client %v", remote.IP)
|
||||
|
||||
var dhcp messages.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)
|
||||
|
||||
options := parseOptions(optData)
|
||||
dhcpMsgType := options[messages.OptionDHCPMessageType]
|
||||
s.Info("DHCP message type ", dhcpMsgType)
|
||||
|
||||
switch messages.MessageType(dhcpMsgType.Data[1]) {
|
||||
case messages.MessageTypeDiscover:
|
||||
go handleOffer(dhcp, remote)
|
||||
|
||||
case messages.MessageTypeRequest:
|
||||
go handleAck(dhcp, remote)
|
||||
}
|
||||
}
|
||||
|
||||
func runServer(interfaceName string) {
|
||||
var addr *net.UDPAddr
|
||||
if interfaceName == "all" {
|
||||
zap.L().Debug("Listening on all interfaces")
|
||||
addr, _ = net.ResolveUDPAddr("udp4", ":67")
|
||||
} else {
|
||||
zap.S().Debugf("Listening on interface %s", interfaceName)
|
||||
iface, err := net.InterfaceByName(interfaceName)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
addrs, err := iface.Addrs()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
// Find first IPv4 address associated with the interface
|
||||
var ip net.IP = nil
|
||||
for _, ifaddr := range addrs {
|
||||
switch a := ifaddr.(type) {
|
||||
case *net.IPAddr:
|
||||
ip = a.IP
|
||||
case *net.IPNet:
|
||||
ip = a.IP
|
||||
}
|
||||
|
||||
ip = ip.To4()
|
||||
if ip != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
if ip == nil {
|
||||
zap.S().Panicf("No IPv4 address associated with interface %s", interfaceName)
|
||||
}
|
||||
fmt.Printf("%+v\n", ip)
|
||||
addr, _ = net.ResolveUDPAddr("udp4", ip.String()+":67")
|
||||
}
|
||||
conn, err := net.ListenUDP("udp4", addr)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
zap.S().Infof("Listening for incoming connections on %s", addr.String())
|
||||
|
||||
for {
|
||||
buf := make([]byte, 1024)
|
||||
rlen, remote, err := conn.ReadFrom(buf[:])
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
// Do stuff with the read bytes
|
||||
remoteAddr, ok := remote.(*net.UDPAddr)
|
||||
if !ok {
|
||||
zap.S().Warn("Not a valid remote IP address: ", remote)
|
||||
continue
|
||||
}
|
||||
go handleMsg(buf[0:rlen], remoteAddr)
|
||||
}
|
||||
}
|
||||
|
||||
var rootCmd = &cobra.Command{
|
||||
Use: "go-dhcp",
|
||||
Short: "go-dhcp is a simple DHCP server written in Go",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
runServer(listenInterface)
|
||||
},
|
||||
}
|
||||
|
||||
func init() {
|
||||
// Set up logging
|
||||
logger, _ := zap.NewDevelopment()
|
||||
defer logger.Sync() // flushes buffer, if any
|
||||
zap.ReplaceGlobals(logger)
|
||||
|
||||
rootCmd.PersistentFlags().StringVarP(&listenInterface, "interface", "i", "all", "Interface to listen on")
|
||||
}
|
||||
|
||||
func Execute() {
|
||||
if err := rootCmd.Execute(); err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
11
cmd/util.go
Normal file
11
cmd/util.go
Normal file
@@ -0,0 +1,11 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
)
|
||||
|
||||
func u32tob(u uint32) []byte {
|
||||
buf := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(buf, u)
|
||||
return buf
|
||||
}
|
||||
Reference in New Issue
Block a user