package main

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"github.com/TransX/log"
	"github.com/TransX/tscipher"
	"net"
	"os"
	"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.Panic("Read panic. Tunnel id: %s. Remote Add: %s. Err:%s", id, src.RemoteAddr().String(), err)
		}
		log.Info("Reived %d bytes. Tunnel: id %s", nByte,id)
		log.Debug("Reived %s",cache[:nByte])
		sendCarrier := &tscipher.Carrier{
			dest,
			tscipher.NewCipher("XOR"),
			cache, //TODO:危险，cache的容量容易被不小心修改
		}
		_, err = tscipher.SendData(sendCarrier, nByte)
		log.Info("Write %d bytes. Tunnel: %s",nByte,id)
		if err != nil {
			log.Panic("Write panic. ID: %s, Err: %s, Remote Add: %s", 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.Panic("Failed to create listener. %s", err)
		os.Exit(0)
	}
	for {
		if listenerConn, err := listener.Accept(); err == nil {
			log.Info("Incoming %s", listenerConn.RemoteAddr().String())
			//创建到目标的连接
			destConn, err := this.createTCPClient(destIP, destPort)
			if err != nil {
				log.Panic("Failed to connect to destination. %s", err)
				os.Exit(0)
			}
			log.Info("Dial %s", destConn.RemoteAddr().String())
			go this.tunnel(listenerConn, destConn, this.tunnelID())
			go this.tunnel(destConn, listenerConn, this.tunnelID())
		} else {
			log.Info("Failed to accept incoming connection. %s", err)
		}
	}
}