go_lang

package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"syscall"
)

type Tunnel struct {
	localAddr  string
	remoteAddr string
}

func NewTunnel(localAddr, remoteAddr string) *Tunnel {
	return &Tunnel{
		localAddr:  localAddr,
		remoteAddr: remoteAddr,
	}
}

func (t *Tunnel) Start() error {
	listener, err := net.Listen("tcp", t.localAddr)
	if err != nil {
		return fmt.Errorf("failed to start listener: %v", err)
	}
	defer listener.Close()

	log.Printf("TCP Tunnel started: %s -> %s\n", t.localAddr, t.remoteAddr)

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		<-sigChan
		log.Println("Shutting down tunnel...")
		listener.Close()
		os.Exit(0)
	}()

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Printf("Error accepting connection: %v\n", err)
			continue
		}

		go t.handleConnection(conn)
	}
}

func (t *Tunnel) handleConnection(localConn net.Conn) {
	defer localConn.Close()

	remoteConn, err := net.Dial("tcp", t.remoteAddr)
	if err != nil {
		log.Printf("Error connecting to remote server: %v\n", err)
		return
	}
	defer remoteConn.Close()

	go func() {
		if _, err := io.Copy(remoteConn, localConn); err != nil {
			log.Printf("Error copying from local to remote: %v\n", err)
		}
	}()

	if _, err := io.Copy(localConn, remoteConn); err != nil {
		log.Printf("Error copying from remote to local: %v\n", err)
	}
}

func main() {
	if len(os.Args) != 3 {
		fmt.Println("Usage: athena <local_addr> <remote_addr>")
		fmt.Println("Example: athena :8080 dhananjaay.dev:80")
		os.Exit(1)
	}

	localAddr := os.Args[1]
	remoteAddr := os.Args[2]

	tunnel := NewTunnel(localAddr, remoteAddr)
	if err := tunnel.Start(); err != nil {
		log.Fatalf("Failed to start tunnel: %v", err)
	}
}