package main

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

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", "", "IP before change")
	newip := flag.String("newip", "10.0.0.15", "IP after change")
	oldmac := flag.String("oldmac", "", "MAC before change")
	newmac := flag.String("newmac", "", "MAC after change")
	passthrough := flag.Bool("passthrough", false, "Whether to pass every traffic through")
	proxy := flag.String("proxy", "1", "Number of the proxy switch")
	pidfile := flag.String("pidfile", "", "Location to write the pid to")
	logfile := flag.String("logfile", "", "Location to write output to")
	flag.Parse()
	log.SetLevel(log.Level(*logLvl))
	OldMac, _ = net.ParseMAC(*oldmac)
	NewMac = util.GenerateMac(*newmac)
	OldIP = net.ParseIP(*oldip).To4()
	NewIP = net.ParseIP(*newip).To4()
	log.SetFormatter(&log.TextFormatter{
		DisableTimestamp: true,
	})
	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)
		}
	}
	util.WritePIDFile(*pidfile)
	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.
// 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)
		isInteresting := false // Debug Help

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

		// Handle IPv6 packet
		if ipv4layer := packet.Layer(layers.LayerTypeIPv4); ipv4layer != nil {
			ipv4Packet, _ := ipv4layer.(*layers.IPv4)
			log.Debug("IP Protocol ", ipv4Packet.Protocol)
			filterIP(prefix, &ipv4Packet.DstIP, &ipv4Packet.SrcIP, 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 DHCP packet (based on IPv4) - drop for now
			if dhcpLayer := packet.Layer(layers.LayerTypeDHCPv4); dhcpLayer != nil && !passthrough {
				//dhcpPacket, _ := dhcpLayer.(*layers.DHCPv4)
				log.Info(prefix, "DHCP packet dropped")
				continue
			}

			// 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 checksum", err)
				}
			}
		}

		// Drop IPv6 packets
		if ipv6layer := packet.Layer(layers.LayerTypeIPv6); ipv6layer != nil && !passthrough {
			log.Info(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.DstProtAddress, &arpPacket.SourceProtAddress, arpPacket.LayerType())
			filterMAC(prefix, &arpPacket.DstHwAddress, &arpPacket.SourceHwAddress, arpPacket.LayerType())
		}

		log.Debug("End packet")

		// Forward original frame to other plug
		if passthrough {
			writer.Write(frameLength)
			writer.Write(frameBytes)
			continue
		}

		// 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)))

		// Write interesting things to debug file
		if isInteresting {
			util.WriteBinary(fmt.Sprintf("/tmp/pck_%di.dat", time.Now().Unix()), frameBytes)
			util.WriteBinary(fmt.Sprintf("/tmp/pck_%do.dat", time.Now().Unix()), newFrameBytes)
			//util.WritePcap("xyz.pcap", packet.Data(), packet.Metadata().CaptureInfo)
		}

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

// filterIP checks whether an IP target selected from src and dst equals a given value. If yes, it is changed
func filterIP(prefix string, dst interface{}, src interface{}, context gopacket.LayerType) {
	var target interface{}
	var condVal net.IP
	var newVal net.IP
	var which string
	if prefix == cmd.In {
		target = dst
		which = "dst"
		condVal = NewIP
		newVal = OldIP
	} else if prefix == cmd.Out {
		target = src
		which = "src"
		condVal = OldIP
		newVal = NewIP
	}
	ip, isIp := target.(*net.IP)
	bs, isBs := target.(*[]byte)

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

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

// filterMAC checks whether a MAC target selected from src and dst equals a given value. If yes, it is changed
func filterMAC(prefix string, dst interface{}, src interface{}, 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 == cmd.In {
			return
		} else if prefix == cmd.Out {
			OldMac = *src.(*net.HardwareAddr)
			log.Info("OldMac set to ", OldMac)
		}
	}

	var target interface{}
	var condVal net.HardwareAddr
	var newVal net.HardwareAddr
	var which string
	if prefix == cmd.In {
		target = dst
		which = "dst"
		condVal = NewMac
		newVal = OldMac
	} else if prefix == cmd.Out {
		target = src
		which = "src"
		condVal = OldMac
		newVal = NewMac
	}
	mac, isMac := target.(*net.HardwareAddr)
	bs, isBs := target.(*[]byte)
	if isMac && bytes.Equal(*mac, condVal) {
		*mac = newVal
		log.Debugf("%s%s %s MAC %s changed to %s", prefix, context, which, condVal, newVal)
	}
	if isBs && bytes.Equal(*bs, condVal) {
		*bs = newVal
		log.Debugf("%s%s %s MAC %s changed to %s", prefix, context, which, condVal, newVal)
	}
}