package main

import (
	"bytes"
	"encoding/binary"
	"flag"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/krolaw/dhcp4"
	log "github.com/sirupsen/logrus"
	"io"
	"net"
	"os"
	"proxy/cmd"
	. "proxy/util"
	"strconv"
	"time"
)

var OldMAC net.HardwareAddr
var NewMAC net.HardwareAddr
var OldIP net.IP
var NewIP net.IP

// In-/Output to VM and network

var VmReader io.Reader
var VmWriter io.Writer
var NetReader io.Reader
var NetWriter io.Writer

// DHCP variables

var DHCPXId []byte
var DHCPIP net.IP
var DHCPRouterIP net.IP
var DHCPDNSIP net.IP
var DHCPMAC net.HardwareAddr
var DHCPMask []byte
var DHCPState dhcp4.MessageType
var DHCPCandidate net.IP

// User-provided variables

var Hostname string
var Passthrough bool
var Wireshark bool

// UId used as host name
var UId string

// Start the two plugs and run two concurrent forward methods
func main() {
	// Get command line arguments
	logLvl := flag.Int("log", 4, "allowed: 5 (debug), 4 (info), 3 (warning), 2 (error), 1 (fatal)")
	oldIP := flag.String("oldip", "", "Force IP before change (optional)")
	newIP := flag.String("newip", "", "Force IP after change (optional)")
	oldMAC := flag.String("oldmac", "", "Force MAC before change (optional)")
	newMAC := flag.String("newmac", "", "Force MAC after change (optional)")
	passthrough := flag.Bool("passthrough", false, "Whether to pass every traffic through")
	sockMain := flag.String("smain", "/run/vde/sw_main.sock", "Main switch sock path, - for stdin/out")
	sockProxy := flag.String("sproxy", "/run/vde/sw_proxy1.sock", "Proxy switch sock path")
	pidFile := flag.String("pidfile", "", "Location to write the pid to (optional)")
	logFile := flag.String("logfile", "", "Location to write output to (optional)")
	wireshark := flag.Bool("wireshark", false, "Whether to write all traffic to /tmp")
	hostname := flag.String("hostname", "", "Set a windows hostname to filter for in binary payloads")
	flag.Parse()
	log.SetLevel(log.Level(*logLvl))
	OldMAC, _ = net.ParseMAC(*oldMAC)
	NewMAC = GenerateMac(*newMAC)
	OldIP = net.ParseIP(*oldIP).To4()
	NewIP = net.ParseIP(*newIP).To4()
	Passthrough = *passthrough
	Wireshark = *wireshark
	Hostname = *hostname
	UId = GenerateUId(*sockProxy)
	log.SetFormatter(&log.TextFormatter{
		DisableTimestamp: true,
	})
	log.Info("Using replacement hostname " + UId)
	log.Info("Using generated MAC " + NewMAC.String())
	if *logFile != "" {
		if f, err := os.OpenFile(*logFile, os.O_WRONLY|os.O_CREATE, 0755); err != nil {
			log.Error("Error opening logFile ", *logFile)
		} else {
			log.SetOutput(f)
		}
	}
	WritePIDFile(*pidFile)
	var c1, c2 *cmd.Cmd
	if *sockMain != "-" {
		c1, NetReader, NetWriter = cmd.Start(*sockMain)
	} else {
		NetReader = os.Stdout
		NetWriter = os.Stdin
	}
	c2, VmReader, VmWriter = cmd.Start(*sockProxy)
	go pipeForward(In)
	go pipeForward(Out)
	if NewIP == nil && !Passthrough {
		sendDHCPRequest(dhcp4.Discover, net.IPv4zero)
	}
	if *sockMain != "-" {
		c1.WaitH()
	}
	c2.WaitH()
}

// Reads from an input and writes to and output,
// do things to the content in between.
// Is meant to be run concurrently with "go pipeForward(...)".
// This is the main loop of the proxy.
func pipeForward(prefix string) {
	var reader io.Reader
	var writer io.Writer
	if prefix == In {
		reader = NetReader
		writer = VmWriter
	} else {
		reader = VmReader
		writer = NetWriter
	}

	for {
		// Read frame length
		frameLength := make([]byte, 2)
		if _, err := reader.Read(frameLength); err == io.EOF {
			log.Fatal(prefix, "Error reading frame length")
		}

		// Read actual frame
		frameBytes := make([]byte, int(binary.BigEndian.Uint16(frameLength)))
		if _, err := reader.Read(frameBytes); err == io.EOF {
			log.Fatal(prefix, "Error reading frame data")
		}

		if Wireshark {
			WritePcap("/tmp/pkg_"+strconv.FormatInt(time.Now().Unix(), 10)+".pcap", frameBytes)
		}

		// Forward original frame to other plug
		if Passthrough {
			if _, err := writer.Write(frameLength); err != nil {
				log.Error("Error forwarding original packet length", err)
			}
			if _, err := writer.Write(frameBytes); err != nil {
				log.Error("Error forwarding original packet data", err)
			}
			continue
		}

		// Convert frame to full stack packet
		packet := gopacket.NewPacket(frameBytes, layers.LayerTypeEthernet, gopacket.Default)

		// Handle Ethernet frame
		log.Debug("Start packet")
		frame := packet.Layer(layers.LayerTypeEthernet).(*layers.Ethernet)
		filterMAC(prefix, &frame.SrcMAC, "src", frame.LayerType())
		filterMAC(prefix, &frame.DstMAC, "dst", frame.LayerType())

		// Handle IPv4 packet
		if ipv4layer := packet.Layer(layers.LayerTypeIPv4); ipv4layer != nil {
			ipv4Packet, _ := ipv4layer.(*layers.IPv4)
			log.Debug("IP Protocol ", ipv4Packet.Protocol)

			// Handle DHCPv4 packet (based on IPv4)
			if dhcpLayer := packet.Layer(layers.LayerTypeDHCPv4); dhcpLayer != nil {
				handleDHCP(dhcpLayer.LayerContents(), frame.DstMAC, frame.SrcMAC, prefix)
				continue
			}

			filterIP(prefix, &ipv4Packet.SrcIP, "src", ipv4Packet.LayerType())
			filterIP(prefix, &ipv4Packet.DstIP, "dst", ipv4Packet.LayerType())

			// Handle ICMP packet (based on IPv4)
			if icmpLayer := packet.Layer(layers.LayerTypeICMPv4); icmpLayer != nil {
				icmpPacket, _ := icmpLayer.(*layers.ICMPv4)
				log.Debug(prefix, "ICMP Type ", icmpPacket.TypeCode)
			}

			// Handle TCP packet
			if tcpLayer := packet.Layer(layers.LayerTypeTCP); tcpLayer != nil {
				tcpPacket, _ := tcpLayer.(*layers.TCP)
				if err := tcpPacket.SetNetworkLayerForChecksum(ipv4Packet); err != nil {
					log.Error(prefix, "Error setting network layer for TCP checksum", err)
				}
			}

			// Handle UDP packet
			if udpLayer := packet.Layer(layers.LayerTypeUDP); udpLayer != nil {
				udpPacket, _ := udpLayer.(*layers.UDP)
				if err := udpPacket.SetNetworkLayerForChecksum(ipv4Packet); err != nil {
					log.Error(prefix, "Error setting network layer for UDP checksum", err)
				}

				if (udpPacket.SrcPort == 137 && udpPacket.DstPort == 137) || // NBNS
					(udpPacket.SrcPort == 138 && udpPacket.DstPort == 138) { // NBDS
					log.Debug(prefix, "Filtering NBNS/NBDS payload")
					filterPayload(prefix, &udpPacket.Payload)
				}
			}
		}

		// Drop IPv6 packets
		if ipv6layer := packet.Layer(layers.LayerTypeIPv6); ipv6layer != nil {
			log.Debug(prefix, "IPv6 packet dropped")
			continue
		}

		// Handle ARP packet
		if frame.EthernetType == layers.EthernetTypeARP {
			arpPacket := packet.Layer(layers.LayerTypeARP).(*layers.ARP)
			log.Debug(prefix, "ARP Type ", arpPacket.Operation)
			filterIP(prefix, &arpPacket.SourceProtAddress, "src", arpPacket.LayerType())
			filterIP(prefix, &arpPacket.DstProtAddress, "dst", arpPacket.LayerType())
			filterMAC(prefix, &arpPacket.SourceHwAddress, "src", arpPacket.LayerType())
			filterMAC(prefix, &arpPacket.DstHwAddress, "dst", arpPacket.LayerType())
		}

		log.Debug("End packet")

		// Serialize packet back to binary
		buf := gopacket.NewSerializeBuffer()
		opts := gopacket.SerializeOptions{ComputeChecksums: true, FixLengths: true}
		if err := gopacket.SerializePacket(buf, opts, packet); err != nil {
			log.Errorf("%s Error serializing packet to send\n%s\nSrc:\t%s\nDst:\t%s\n", prefix, err, frame.SrcMAC, frame.DstMAC)
			continue
		}
		newFrameBytes := buf.Bytes()
		newFrameLength := make([]byte, 2)
		binary.BigEndian.PutUint16(newFrameLength, uint16(len(newFrameBytes)))

		// Forward modified frame to other plug
		if _, err := writer.Write(newFrameLength); err != nil {
			log.Error("Error forwarding packet length", err)
		}
		if _, err := writer.Write(newFrameBytes); err != nil {
			log.Error("Error forwarding packet data", err)
		}
	}
}

// filterPayloads filters binary payloads to be able to process unsupported protocols
func filterPayload(prefix string, payload *[]byte) {
	// Populate slices with values to on vm and network side
	vmVals := [][]byte{OldMAC}
	netVals := [][]byte{NewMAC}
	if OldIP != nil && !OldIP.Equal(NewIP) {
		vmVals = append(vmVals, OldIP)
		netVals = append(netVals, NewIP)
	}
	if Hostname != "" {
		vmVals = append(vmVals, []byte(Hostname), DnsQueryEncode(Hostname))
		netVals = append(netVals, []byte(UId), DnsQueryEncode(UId))
	}

	// Choose for which slice to search for and which to replace it with
	searchVals := vmVals
	replaceVals := netVals
	if prefix == In {
		searchVals = netVals
		replaceVals = vmVals
	}

	// Replace all occurencies of each slice value
	for i, search := range searchVals {
		*payload = bytes.Replace(*payload, search, replaceVals[i], -1)
	}
}

// filterIP checks whether an IP target selected from src and dst equals a given value. If yes, it is changed
func filterIP(prefix string, addr interface{}, which string, context gopacket.LayerType) {
	ip, isIp := addr.(*net.IP)
	bs, isBs := addr.(*[]byte)

	// If no OldIP is set yet, get it from outgoing src field
	if OldIP == nil {
		if prefix == In || which == "dst" {
			return
		} else if prefix == Out && which == "src" {
			if isIp {
				if !ip.IsGlobalUnicast() {
					return
				}
				OldIP = *ip
			} else if isBs {
				OldIP = *bs
			}
			log.Info("OldIP set to ", OldIP)
		}
	}

	if isIp && bytes.Equal(*ip, OldIP) {
		*ip = NewIP
		log.Debugf("%s%s %s IP %s changed to %s", prefix, context, which, OldIP, NewIP)
	} else if isBs && bytes.Equal(*bs, OldIP) {
		*bs = NewIP
		log.Debugf("%s%s %s IP %s changed to %s", prefix, context, which, OldIP, NewIP)
	} else if isIp && bytes.Equal(*ip, NewIP) {
		*ip = OldIP
		log.Debugf("%s%s %s IP %s changed to %s", prefix, context, which, NewIP, OldIP)
	} else if isBs && bytes.Equal(*bs, NewIP) {
		*bs = OldIP
		log.Debugf("%s%s %s IP %s changed to %s", prefix, context, which, NewIP, OldIP)
	}
}

// filterMAC checks whether a MAC target selected from src and dst equals a given value. If yes, it is changed
func filterMAC(prefix string, addr interface{}, which string, context gopacket.LayerType) {
	// If no OldMAC is set yet, get it from outgoing src field
	// Has to be HardwareAddr because this is used for ethernet frames which call this method first
	if OldMAC == nil {
		if prefix == In || which == "dst" {
			return
		} else if prefix == Out && which == "src" {
			OldMAC = *addr.(*net.HardwareAddr)
			log.Info("OldMAC set to ", OldMAC)
		}
	}
	mac, isMac := addr.(*net.HardwareAddr)
	bs, isBs := addr.(*[]byte)
	if isMac && bytes.Equal(*mac, NewMAC) {
		*mac = OldMAC
		log.Debugf("%s%s %s MAC %s changed to %s", prefix, context, which, NewMAC, OldMAC)
	} else if isBs && bytes.Equal(*bs, NewMAC) {
		*bs = OldMAC
		log.Debugf("%s%s %s MAC %s changed to %s", prefix, context, which, NewMAC, OldMAC)
	} else if isMac && bytes.Equal(*mac, OldMAC) {
		*mac = NewMAC
		log.Debugf("%s%s %s MAC %s changed to %s", prefix, context, which, OldMAC, NewMAC)
	} else if isBs && bytes.Equal(*bs, OldMAC) {
		*bs = NewMAC
		log.Debugf("%s%s %s MAC %s changed to %s", prefix, context, which, OldMAC, NewMAC)
	}
}

// handleDHCP provides the main DHCP functionality to request an IP from a server on the one hand
// and provide an IP to the VM on the other hand
func handleDHCP(content []byte, dstMAC net.HardwareAddr, srcMAC net.HardwareAddr, prefix string) {
	req := dhcp4.Packet(content)
	if req.HLen() > 16 { // Invalid size
		log.Error(prefix, "Invalid DHCP size")
		return
	}
	options := req.ParseOptions()
	var reqType dhcp4.MessageType
	if t := options[dhcp4.OptionDHCPMessageType]; len(t) != 1 {
		log.Error(prefix, "Invalid DHCP message type")
		return
	} else {
		reqType = dhcp4.MessageType(t[0])
		if reqType < dhcp4.Discover || reqType > dhcp4.Inform {
			log.Error(prefix, "Invalid DHCP message type: ", reqType)
			return
		}
	}
	log.Debug(prefix, "DHCP message registered: ", reqType)
	if prefix == Out {
		switch reqType {
		case dhcp4.Discover:
			if OldIP == nil {
				if NewIP == nil || !NewIP.IsGlobalUnicast() {
					log.Warning(prefix, "DHCP request, but neither OldIP nor NewIP known")
					sendDHCPRequest(dhcp4.Discover, net.IPv4zero)
					return
				}
				sendDHCPReply(req, dhcp4.Offer, NewIP, options, dstMAC)
			} else {
				sendDHCPReply(req, dhcp4.Offer, OldIP, options, dstMAC)
			}
		case dhcp4.Inform:
			fallthrough
		case dhcp4.Request:
			reqIP := net.IP(options[dhcp4.OptionRequestedIPAddress])
			if reqIP == nil {
				reqIP = req.CIAddr()
			}
			if len(reqIP) != 4 || reqIP.Equal(net.IPv4zero) {
				log.Error(prefix, "Invalid IP requested in DHCP: ", reqIP)
				return
			}
			sendDHCPReply(req, dhcp4.ACK, reqIP, options, dstMAC)
		}
	} else {
		switch reqType {
		case dhcp4.Offer:
			if DHCPState == dhcp4.Discover {
				DHCPMAC = srcMAC
				offIP := req.YIAddr()
				if len(offIP) != 4 || offIP.Equal(net.IPv4zero) {
					log.Error(prefix, "Invalid IP offered in DHCP: ", offIP)
					return
				}
				if dhcpip := options[dhcp4.OptionServerIdentifier]; dhcpip != nil {
					DHCPIP = dhcpip
				}
				if mask := options[dhcp4.OptionSubnetMask]; mask != nil {
					DHCPMask = mask
				}
				if dns := options[dhcp4.OptionDomainNameServer]; dns != nil {
					DHCPDNSIP = dns
				}
				if router := options[dhcp4.OptionRouter]; router != nil {
					DHCPRouterIP = router
				}
				DHCPCandidate = offIP
				sendDHCPRequest(dhcp4.Request, offIP)
			}
		case dhcp4.ACK:
			if DHCPState == dhcp4.Request {
				NewIP = DHCPCandidate
				DHCPCandidate = nil
				DHCPState = 0
				log.Info("DHCP lease accepted: ", NewIP)
			}
		case dhcp4.NAK:
			if DHCPState == dhcp4.Request {
				sendDHCPRequest(dhcp4.Discover, net.IPv4zero)
			}
		}
	}
}

// sendDHCPReply creates a response DHCP packet for the VM and sends it
func sendDHCPReply(req dhcp4.Packet, mt dhcp4.MessageType, lease net.IP, reqOpt dhcp4.Options, dstMAC net.HardwareAddr) {
	if DHCPIP == nil || DHCPMAC == nil {
		log.Info("DHCP server is not known, discover request from VM discarded")
		sendDHCPRequest(dhcp4.Discover, net.IPv4zero)
		return
	}

	log.Info("Sending DHCP response: ", mt, lease)

	// Getting the options
	opt := dhcp4.Options{
		dhcp4.OptionSubnetMask:       DHCPMask,
		dhcp4.OptionRouter:           DHCPRouterIP,
		dhcp4.OptionDomainNameServer: DHCPDNSIP,
	}.SelectOrderOrAll(reqOpt[dhcp4.OptionParameterRequestList])

	// Creating the full packet layer by layer
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		ComputeChecksums: true,
		FixLengths:       true,
	}
	eth := layers.Ethernet{
		SrcMAC:       DHCPMAC,
		DstMAC:       dstMAC,
		EthernetType: layers.EthernetTypeIPv4,
	}
	ipv4 := layers.IPv4{
		Version:  4,
		TTL:      128,
		Protocol: layers.IPProtocolUDP,
		SrcIP:    DHCPIP,
		DstIP:    net.IPv4bcast,
	}
	udp := layers.UDP{
		SrcPort: 67,
		DstPort: 68,
	}
	dhcp := dhcp4.ReplyPacket(req, mt, DHCPIP, lease, 24*time.Hour, opt)
	if err := udp.SetNetworkLayerForChecksum(&ipv4); err != nil {
		log.Error("Error building DHCP response:", err)
	}
	if err := gopacket.SerializeLayers(buf, opts, &eth, &ipv4, &udp, gopacket.Payload(dhcp)); err != nil {
		log.Error("Error serializing DHCP response:", err)
	}
	packetData := buf.Bytes()

	// Sending layer through VM's pipe
	packetLength := make([]byte, 2)
	binary.BigEndian.PutUint16(packetLength, uint16(len(packetData)))

	if _, err := VmWriter.Write(packetLength); err != nil {
		log.Error("Error writing DHCP response length", err)
	}
	if _, err := VmWriter.Write(packetData); err != nil {
		log.Error("Error writing DHCP response data", err)
	}
}

// sendDHCPRequest creates a request DHCP packet for the server and sends it
func sendDHCPRequest(mt dhcp4.MessageType, reqIP net.IP) {
	log.Info("Sending DHCP request: ", mt)
	if mt == dhcp4.Discover {
		DHCPXId = GenerateXID()
	}
	DHCPState = mt

	// Creating the full packet layer by layer
	buf := gopacket.NewSerializeBuffer()
	serializeOpts := gopacket.SerializeOptions{
		ComputeChecksums: true,
		FixLengths:       true,
	}
	eth := layers.Ethernet{
		SrcMAC:       NewMAC,
		DstMAC:       If(mt == dhcp4.Discover).MAC([]byte{255, 255, 255, 255, 255, 255}, DHCPMAC),
		EthernetType: layers.EthernetTypeIPv4,
	}
	ipv4 := layers.IPv4{
		Version:  4,
		TTL:      128,
		Protocol: layers.IPProtocolUDP,
		SrcIP:    net.IPv4zero,
		DstIP:    net.IPv4bcast,
	}
	udp := layers.UDP{
		SrcPort: 68,
		DstPort: 67,
	}
	dhcpOpts := []dhcp4.Option{
		{
			Code:  dhcp4.OptionHostName,
			Value: []byte("vdeproxy" + UId),
		},
		{
			Code:  dhcp4.OptionParameterRequestList,
			Value: []byte{1, 3, 6}, // Subnet Mask, Router, Domain Name Server
		},
	}
	if mt == dhcp4.Request {
		dhcpOpts = append(dhcpOpts, dhcp4.Option{
			Code:  dhcp4.OptionRequestedIPAddress,
			Value: reqIP.To4(),
		})
	}
	dhcp := dhcp4.RequestPacket(mt, NewMAC, reqIP, DHCPXId, false, dhcpOpts)
	if err := udp.SetNetworkLayerForChecksum(&ipv4); err != nil {
		log.Error("Error building DHCP request: ", err)
	}
	if err := gopacket.SerializeLayers(buf, serializeOpts, &eth, &ipv4, &udp, gopacket.Payload(dhcp)); err != nil {
		log.Error("Error serializing DHCP request: ", err)
	}
	packetData := buf.Bytes()

	// Sending layer through VM's pipe
	packetLength := make([]byte, 2)
	binary.BigEndian.PutUint16(packetLength, uint16(len(packetData)))
	if _, err := NetWriter.Write(packetLength); err != nil {
		log.Error("Error writing DHCP response length: ", err)
	}
	if _, err := NetWriter.Write(packetData); err != nil {
		log.Error("Error writing DHCP response data: ", err)
	}
}