package main

import (
	"flag"
	"io"
	"net"
	"os/exec"
	"time"

	"github.com/phuslu/log"
)

func forward(preCommand string, lconn net.Conn, raddr string) {
	if preCommand != "" {
		if output, err := exec.Command("/bin/bash", "-c", preCommand).CombinedOutput(); err != nil {
			log.Error().Err(err).Str("pre_command", preCommand).Bytes("output", output).Msg("exec pre_command error")
		} else {
			log.Error().Str("pre_command", preCommand).Bytes("output", output).Msg("exec pre_command ok")
		}
	}

	log.Info().Str("remote_addr", raddr).Stringer("local_addr", lconn.LocalAddr()).Msg("connecting to remote address")
	rconn, err := net.Dial("tcp", raddr)
	if err != nil {
		log.Error().Err(err).Str("remote_addr", raddr).Msg("dail remote address failed")
		return
	}

	log.Info().Stringer("remote_addr", rconn.RemoteAddr()).Stringer("local_addr", lconn.LocalAddr()).Msg("forward")
	go func() {
		defer rconn.Close()
		defer lconn.Close()
		io.Copy(rconn, lconn)
	}()
	go func() {
		defer lconn.Close()
		defer rconn.Close()
		io.Copy(lconn, rconn)
	}()
}

func main() {
	var laddr, raddr, command string
	flag.StringVar(&laddr, "laddr", "0.0.0.0:2222", "local address")
	flag.StringVar(&raddr, "raddr", "192.168.1.3:2222", "remote address")
	flag.StringVar(&command, "command", "", "pre command")

	flag.Parse()

	ln, err := net.Listen("tcp", laddr)
	if err != nil {
		log.Fatal().Str("local_addr", laddr).Msg("listen address failed")
	}
	defer ln.Close()

	log.Info().Stringer("local_addr", ln.Addr()).Msg("serving...")
	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Error().Err(err).Msg("accpet connection error")
			time.Sleep(10 * time.Millisecond)
			continue
		}
		go forward(command, conn, raddr)
	}
}