package main

import (
	"encoding/binary"
	"flag"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	log "github.com/sirupsen/logrus"
	"io"
	"net"
	"proxy/cmd"
	"strconv"
)

var OldMac net.HardwareAddr
var NewMac net.HardwareAddr
var OldIP net.IP
var NewIP net.IP


// 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", "10.0.0.11", "IP before change")
	newip := flag.String("newip", "10.0.0.15", "IP after change")
	oldmac := flag.String("oldmac", "52:54:00:12:34:56", "MAC before change")
	newmac := flag.String("newmac", "52:54:00:12:34:aa", "MAC after change")
	passthrough := flag.Bool("passthrough", false, "Whether to pass every traffic through")
	proxy := flag.String("proxy", "1", "Number of the proxy switch")
	flag.Parse()
	log.SetLevel(log.Level(*logLvl))
	OldMac, _ = net.ParseMAC(*oldmac)
	NewMac, _ = net.ParseMAC(*newmac)
	OldIP = net.ParseIP(*oldip)
	NewIP = net.ParseIP(*newip)
	log.SetFormatter(&log.TextFormatter{
		DisableTimestamp: true,
	})
	c1 := cmd.New("vde_plug", "/run/vde/sw_main.sock")
	c2 := cmd.New("vde_plug", "/run/vde/sw_proxy" + *proxy + ".sock")
	c1.Execute()
	c2.Execute()
	go pipeForward(c1.OutReader, c2.InWriter, cmd.In, *passthrough)
	go pipeForward(c2.OutReader, c1.InWriter, cmd.Out, *passthrough)
	c1.WaitH()
	c2.WaitH()
}

// Reads from an input and writes to and output,
// do things to the content in between.
// For now only output the packet's information
// Is meant to be run concurrently with "go pipeForward(...)"
func pipeForward(reader io.Reader, writer io.Writer, prefix string, passthrough bool) {
	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")
		}

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

		// Handle Ethernet frame
		frame := packet.Layer(layers.LayerTypeEthernet).(*layers.Ethernet)
		if !passthrough && prefix == cmd.In {
			log.Debug(prefix, "Incoming MAC rewritten ", frame.DstMAC, OldMac)
			frame.DstMAC = OldMac
		} else if !passthrough && prefix == cmd.Out {
			log.Debug(prefix, "Outgoing MAC rewritten ", frame.SrcMAC, NewMac)
			frame.SrcMAC = NewMac
		}

		// Handle IP packet
		if ipv4layer := packet.Layer(layers.LayerTypeIPv4); ipv4layer != nil {
			ipv4Packet, _ := ipv4layer.(*layers.IPv4)
			if !passthrough && prefix == cmd.In {
				log.Debug(prefix, "Incoming IP rewritten ", ipv4Packet.DstIP, OldIP)
				ipv4Packet.DstIP = OldIP
			} else if !passthrough && prefix == cmd.Out {
				log.Debug(prefix, "Outgoing IP rewritten ", ipv4Packet.SrcIP, NewMac)
				ipv4Packet.SrcIP = NewIP
			}
		}

		// Handle ARP packet
		if frame.EthernetType == layers.EthernetTypeARP {
			arpPacket := packet.Layer(layers.LayerTypeARP).(*layers.ARP)
			log.Infof(
				"%sARP before modification\nAddrType:\t%s\nProtocol:\t%s\nOperation:\t%s\n" +
					"SrcHwAddress:\t%s\nSrcProtAddress:\t%s\nDstHwAddress:\t%s\nDstProAddress:\t%s\n",
					prefix,
					arpPacket.AddrType,
					arpPacket.Protocol,
					strconv.Itoa(int(arpPacket.Operation)),
					net.HardwareAddr(arpPacket.SourceHwAddress),
					net.IP(arpPacket.SourceProtAddress),
					net.HardwareAddr(arpPacket.DstHwAddress),
					net.IP(arpPacket.DstProtAddress),
				)

			if !passthrough && prefix == cmd.In {
				log.Debug(prefix,"Incoming MAC rewritten (ARP) ", net.HardwareAddr(arpPacket.DstHwAddress), OldMac)
				arpPacket.DstProtAddress = OldIP
				log.Infof(
					"%sARP after modification\nDstProtAddress:\t%s\n",
					prefix,
					net.IP(arpPacket.DstProtAddress),
				)
				if arpPacket.Operation == layers.ARPReply {
					arpPacket.DstHwAddress = OldMac
					log.Infof(
						"%sDstHwAddress:\t%s\n",
						prefix,
						net.HardwareAddr(arpPacket.DstHwAddress),
					)
				}
			} else if !passthrough && prefix == cmd.Out {
				log.Debug(prefix,"Outgoing MAC rewritten (ARP) ", net.HardwareAddr(arpPacket.SourceHwAddress), NewMac)
				arpPacket.SourceHwAddress = NewMac
				arpPacket.SourceProtAddress = NewIP
				log.Infof(
					"%sARP after modification\nSrcHwAddress:\t%s\nSrcProtAddress:\t%s\n",
					prefix,
					net.HardwareAddr(arpPacket.SourceHwAddress),
					net.IP(arpPacket.SourceProtAddress),
				)
			}
		}

		// Handle DHCP packet
		if dhcpLayer := packet.Layer(layers.LayerTypeDHCPv4); dhcpLayer != nil {
			//dhcpPacket, _ := dhcpLayer.(*layers.DHCPv4)
			log.Info(prefix, "DHCP packet dropped")
			continue
		}

		buf := gopacket.NewSerializeBuffer()
		if err := gopacket.SerializePacket(buf, gopacket.SerializeOptions{}, packet); err != nil {
			log.Error(prefix, "Error serializing packet to send")
			continue
		}
		newFrameBytes := buf.Bytes()
		newFrameLength := make([]byte, 2)
		binary.BigEndian.PutUint16(newFrameLength, uint16(len(newFrameBytes)))

		// Forward original frame to other plug
		writer.Write(newFrameLength)
		writer.Write(newFrameBytes)
	}
}