package shadow

import (
	"fmt"
	"io"
	"net"
	"time"
)

type Server struct {
	ListenAddress string
	TargetAddress string
	FakeAddress   string
}

func NewServer(listenAddress string, targetAddress string, fakeAddress string) *Server {
	server := &Server{
		ListenAddress: listenAddress,
		TargetAddress: targetAddress,
		FakeAddress:   fakeAddress,
	}
	return server
}

func (s *Server) Start() {
	listen, err := net.Listen("tcp", s.ListenAddress)
	if err != nil {
		fmt.Printf("[Server] Start server error: %v\n", err)
		return
	}
	defer listen.Close()
	fmt.Printf("[Server] Listening at:%v\n", s.ListenAddress)
	for {
		conn, err := listen.Accept()
		if err != nil {
			fmt.Printf("[Server] Accept error: %v\n", err)
			continue
		}
		go handler(conn, s.TargetAddress, s.FakeAddress)
	}
}

func handler(conn net.Conn, targetAddress string, fakeAddress string) {

	//Process fake TLS handshake
	fmt.Println("[Server] Perform handshake")
	fakeConn, err := net.Dial("tcp", fakeAddress)
	if err != nil {
		fmt.Printf("[Server] Dial fake error : %v\n", err)
		return
	}
	waitCh := make(chan int, 1)

	go processHandshake(conn, fakeConn, waitCh, "client")
	go processHandshake(fakeConn, conn, waitCh, "server")

	<-waitCh

	//Clean up previous buffered data
	conn.SetDeadline(time.Now())
	conn.SetDeadline(time.Time{})
	fakeConn.Close()

	realConnection, err := net.Dial("tcp", targetAddress)
	if err != nil {
		fmt.Printf("[Server] Dial target error : %v\n", err)
		return
	}
	if err != nil {
		return
	}

	p := &PackAppData{Conn: conn}

	defer p.Close()
	defer realConnection.Close()
	exit := make(chan int, 1)

	go MyCopy(p, realConnection, exit)
	go MyCopy(realConnection, p, exit)
	<-exit
}

func processHandshake(src net.Conn, dst net.Conn, waitCh chan int, srcType string) {
	buf := make([]byte, 32*1024)
	verifyPass := false
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if srcType == "client" {
				header := ParseAndVerifyTLSHeader(buf[0:nr])
				if header != nil {
					if header != nil && header.Type == Handshake && header.HandshakeType == ClientHello && !verifyPass {
						verifyPass = VerifyKey(header.Rand, HandshakePassword)
					}
					if header.Type == ChangeCipherSpec {
						if HandshakePassword != "" && !verifyPass {
							fmt.Println("[Server] Probe detected,pass through all traffic.")
						} else {
							fmt.Println("[Server] handshake complete")
							waitCh <- 1
							break
						}
					}
					//fmt.Println(header.toString())
				}
			}
			if nw < 0 || nr < nw {
				nw = 0
				if ew == nil {
					//fmt.Printf("ERR1 %v \n", ew)
				}
			}
			if ew != nil {
				//fmt.Printf("ERR2 %v \n", ew)
				break
			}
			if nr != nw {
				//fmt.Printf("ERR3 %v \n", "shortwrite")
				break
			}
		}
		if er != nil {
			if er != io.EOF {
				//fmt.Printf("ERR4 %v \n", er)
			}
			break
		}
	}
}