package main

import (
    "io"
    "log"
    "net"
    "os"
    "strconv"
    "time"

    "github.com/google/uuid"
)

func main() {
    listenAddr := ":2222" // The local port on which to listen for incoming SSH connections.
    backendAddr := os.Getenv("SSH_BACKEND") // Backend SSH server address from environment variable.
    maxDuration := os.Getenv("SSH_MAX_DURATION") // Max connection duration from environment variable.

    duration, err := strconv.Atoi(maxDuration)
    if err != nil {
        log.Fatalf("Invalid SSH_MAX_DURATION value: %s", err)
    }

    listener, err := net.Listen("tcp", listenAddr)
    if err != nil {
        log.Fatalf("Failed to open listener: %s", err)
    }
    defer listener.Close()
    log.Println("Listening on", listenAddr)

    for {
        clientConn, err := listener.Accept()
        if err != nil {
            log.Println("Failed to accept connection:", err)
            continue
        }

        go handleConnection(clientConn, backendAddr, time.Duration(duration)*time.Second)
    }
}

func handleConnection(clientConn net.Conn, backendAddr string, maxDuration time.Duration) {
    connID := uuid.New().String()
    log.Printf("New connection [%s] started from %s", connID, clientConn.RemoteAddr())

    defer clientConn.Close()

    backendConn, err := net.Dial("tcp", backendAddr)
    if err != nil {
        log.Printf("Failed to connect to backend [%s]: %s", connID, err)
        return
    }
    defer backendConn.Close()

    // Set up a timer to close both connections when the maxDuration is exceeded
    timer := time.AfterFunc(maxDuration, func() {
        log.Printf("Connection [%s] exceeded max duration, terminating", connID)
        clientConn.Close()
        backendConn.Close()
    })
    defer timer.Stop()

    // Forward traffic between the client and the backend
    go func() {
        _, err := io.Copy(backendConn, clientConn)
        if err != nil {
            log.Printf("Error forwarding from client to backend [%s]: %s", connID, err)
        }
    }()

    _, err = io.Copy(clientConn, backendConn)
    if err != nil {
        log.Printf("Error forwarding from backend to client [%s]: %s", connID, err)
    }

    log.Printf("Connection [%s] terminated", connID)
}