package main

import (
	"flag"
	"fmt"
	"git.eve.moe/jackyyf/transparent-proxy/server"
	"io"
	"net"
)

var port = flag.Int("port", 1080, "Listen port of the server")

func main() {
	flag.Parse()
	go IPv4Handler()
	go IPv6Handler()
	select {}
}

func IPv4Handler() {
	proxy_listener, err := server.NewIPv4TransparentListener(fmt.Sprintf("0.0.0.0:%d", *port))
	if err != nil {
		panic(err)
	}
	for {
		conn, err := proxy_listener.Accept()
		if err != nil {
			fmt.Println("Accept Error: ", err.Error())
			continue
		}
		// fmt.Printf("Accepted connection: %s => %s\n", conn.TCPConn().RemoteAddr().String(), conn.RealAddr().String())
		go handle4(conn)
	}
}

func IPv6Handler() {
	proxy_listener, err := server.NewIPv6TransparentListener(fmt.Sprintf("[::]:%d", *port))
	if err != nil {
		panic(err)
	}
	for {
		conn, err := proxy_listener.Accept()
		if err != nil {
			fmt.Println("Accept Error: ", err.Error())
			continue
		}
		// fmt.Printf("Accepted connection: %s => %s\n", conn.TCPConn().RemoteAddr().String(), conn.RealAddr().String())
		go handle6(conn)
	}
}

func pipeThenClose(reader io.Reader, writer io.WriteCloser) {
	defer writer.Close()
	buff := make([]byte, 4096)
	for {
		n, err := reader.Read(buff)
		if n > 0 {
			_, err := writer.Write(buff[:n])
			if err != nil {
				fmt.Printf("Write Error: %#v\n", err.Error())
				return
			}
		}
		if err == io.EOF {
			// fmt.Println("Reader reached EOF, closing.")
			return
		}
		if err != nil {
			// Enable this line ONLY FOR DEBUG PURPOSE.
			// fmt.Printf("Read Error: %#v\n", err.Error())
			return
		}
	}
}

func handle4(conn server.TransparentConnection) {
	rconn, err := net.DialTCP("tcp4", nil, conn.RealAddr())
	if err != nil {
		fmt.Println("Connect to ", conn.RealAddr().String(), " error: ", err.Error())
		conn.TCPConn().Close()
		return
	}
	// fmt.Printf("ESTABLISHED %s <=> %s\n", conn.TCPConn().RemoteAddr().String(), conn.RealAddr().String())
	go pipeThenClose(conn.TCPConn(), rconn)
	go pipeThenClose(rconn, conn.TCPConn())
}

func handle6(conn server.TransparentConnection) {
	rconn, err := net.DialTCP("tcp6", conn.TCPConn().RemoteAddr().(*net.TCPAddr), conn.RealAddr())
	if err != nil {
		fmt.Println("Connect to ", conn.RealAddr().String(), " error: ", err.Error())
		conn.TCPConn().Close()
		return
	}
	// fmt.Printf("ESTABLISHED %s <=> %s\n", conn.TCPConn().RemoteAddr().String(), conn.RealAddr().String())
	go pipeThenClose(conn.TCPConn(), rconn)
	go pipeThenClose(rconn, conn.TCPConn())
}