package main

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"github.com/TransX/tscipher"
	"log"
	"net"
	"strconv"
	"sync/atomic"
	"time"
)

type TransTCP struct {
	seed int32
}

func NewTransTCP() *TransTCP {

	return &TransTCP{
		0,
	}
}

func (this *TransTCP) createTCPClient(ip, port string) (conn net.Conn, err error) {
	conn, err = net.Dial("tcp4", ip+":"+port)
	if err == nil {

	} else {
		conn = nil
	}
	return
}

func (this *TransTCP) createTCPListener(ip, port string) (listen net.Listener, err error) {
	listener, _err := net.Listen("tcp4", ip+":"+port)
	if _err == nil {
		listen = listener
		err = nil
		return
	} else {
		listen = nil
		err = _err
	}
	return
}

func (this *TransTCP) tunnel(src, dest net.Conn, id string) {
	defer func() {
		if r := recover(); r != nil {
			if src != nil {
				src.Close()
			}
			if dest != nil {
				dest.Close()
			}

		}
	}()
	cache := make([]byte, 1024*128) //128kB
	for {
		//构建Carrier
		revCarrier := &tscipher.Carrier{
			src,
			tscipher.NewCipher("XOR"),
			cache,
		}
		nByte, err := tscipher.ReceiveData(revCarrier)
		if err != nil {
			log.Panicln("Read panic", id, err, src.RemoteAddr().String())
		}
		log.Println("Reived ", nByte, "bytes:", id, string(cache[:nByte]))
		sendCarrier := &tscipher.Carrier{
			dest,
			tscipher.NewCipher("XOR"),
			cache, //TODO:危险，cache的容量容易被不小心修改
		}
		_, err = tscipher.SendData(sendCarrier, nByte)
		log.Println("Write")
		if err != nil {
			log.Panicln("Write panic", id, err, dest.RemoteAddr().String())
		}
	}

}

func (this *TransTCP) tunnelID() string {
	nowString := time.Now().String() + strconv.Itoa(int(this.seed))
	atomic.AddInt32(&this.seed, 1) //避免多线程情况下获得的种子相同
	md5Byte := md5.Sum(bytes.NewBufferString(nowString).Bytes())
	return hex.EncodeToString(md5Byte[:])
}

func (this *TransTCP) Start(listenPort, destIP, destPort string) {

	listener, err := this.createTCPListener("0.0.0.0", listenPort)
	if err != nil {
		log.Fatalln("Failed to create listener.", err)
	}
	for {
		if listenerConn, err := listener.Accept(); err == nil {
			log.Println("Incoming ", listenerConn.RemoteAddr().String())
			//创建到目标的连接
			destConn, err := this.createTCPClient(destIP, destPort)
			if err != nil {
				log.Fatalln("Failed to connect to destination.", err)
			}
			log.Println("Dial", destConn.RemoteAddr().String())
			go this.tunnel(listenerConn, destConn, this.tunnelID())
			go this.tunnel(destConn, listenerConn, this.tunnelID())
		} else {
			log.Println("Failed to accept incoming connection.", err)
		}
	}
}