package main

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

var seed int32

func init() {
	seed = 0
}

type TransTCP struct {
}

func NewTransTCP() *TransTCP {

	return &TransTCP{}
}

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, encrypDirection string) { //单向的，从src发送到dest
	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,
		}
		log.Debug("Encrypt Direction %s ID %s", encrypDirection, id)
		if encrypDirection != "receive" {
			revCarrier.Cipher = nil
			log.Debug("Read not crypted. Tunnel: %s", id)
		}
		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 from %s. Tunnel: id %s", nByte, src.RemoteAddr().String(), id)
		log.Debug("Reived %s %s", id, cache[:nByte])
		sendCarrier := &tscipher.Carrier{
			dest,
			tscipher.NewCipher("XOR"),
			cache, //TODO:危险，cache的容量容易被不小心修改
		}
		if encrypDirection != "send" {
			sendCarrier.Cipher = nil
			log.Debug("Write not crypted. Tunnel: %s", id)
		}
		_, err = tscipher.SendData(sendCarrier, nByte)
		log.Info("Write %d bytes from %s to %s. Tunnel: %s", nByte, dest.LocalAddr(), dest.RemoteAddr().String(), id)
		log.Debug("Write %s %s", id, cache[:nByte])
		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(seed))
	atomic.AddInt32(&seed, 1) //避免多线程情况下获得的种子相同
	md5Byte := md5.Sum(bytes.NewBufferString(nowString).Bytes())
	//	log.Info("seed %d %s", seed, hex.EncodeToString(md5Byte[:]))
	return hex.EncodeToString(md5Byte[:])
}

func (this *TransTCP) Start(listenPort, destIP, destPort string, clientOrServer 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())
			//tunnel model : [ -->>server ---- client -->> ](this is a tunnel)
			tunnelIDA := this.tunnelID()
			tunnelIDB := this.tunnelID()
			if clientOrServer == "client" {
				go this.tunnel(listenerConn, destConn, tunnelIDA, "send")
				go this.tunnel(destConn, listenerConn, tunnelIDB, "receive")
				log.Debug("tow tunnel created: %s %s %s %s", tunnelIDA, "send", tunnelIDB, "receive")
			}
			if clientOrServer == "server" {
				go this.tunnel(listenerConn, destConn, tunnelIDA, "receive")
				go this.tunnel(destConn, listenerConn, tunnelIDB, "send")
				log.Debug("tow tunnel created: %s %s %s %s", tunnelIDA, "receive", tunnelIDB, "send")
			}

		} else {
			log.Info("Failed to accept incoming connection. %s", err)
		}
	}
}