gus/main.go

// Copyright 2017-2025 WIT.COM Inc. All rights reserved.
// Use of this source code is governed by the GPL 3.0

package main

import (
	"bufio"
	"embed"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	sync "sync"
	"time"

	"github.com/google/uuid"
	"go.wit.com/dev/alexflint/arg"
	"go.wit.com/gui"
	"go.wit.com/log"
	timestamppb "google.golang.org/protobuf/types/known/timestamppb"
)

var VERSION string
var BUILDTIME string

//go:embed resources/*
var resources embed.FS

func main() {
	var pp *arg.Parser
	gui.InitArg()
	pp = arg.MustParse(&argv)

	if pp == nil {
		pp.WriteHelp(os.Stdout)
		os.Exit(0)
	}
	log.Info("tmp hack", uuid.New().String())

	me = new(gusconf)
	me.pollDelay = 10 * time.Second
	me.portmaps = ConfigLoad()
	me.events = EventLoad()

	if me.portmaps == nil {
		me.portmaps = NewPortmaps()
		p := new(Portmap)
		p.Dest = "testing:323"
		me.portmaps.Append(p)
	}

	if argv.Daemon {
		// turn off timestamps for STDOUT (systemd adds them)
		log.DaemonMode(true)
		startGus()
		startHTTP()
		os.Exit(0)
	}

	if gui.NoGui() {
		startGus()
		startHTTP()
		os.Exit(0)
	}

	startGus()
	// go NewWatchdog()
	go startHTTP()
	doGui()
}

func startGus() {
	all := me.portmaps.All()
	for all.Scan() {
		pm := all.Next()
		if !pm.Enabled {
			continue
		}
		log.Info("portmap enabled for port", pm.Localport, "to", pm.Dest)
		go gus3000(pm)
	}
}

// func doME(pm *Portmap, gus listener.Accept) {
func doME(pm *Portmap, gus net.Listener) {
	localport := int(pm.Localport)
	where := pm.Dest

	/*
		// Listen on local port 3000
		s := fmt.Sprintf("0.0.0.0:%d", port)
		listener, err := net.Listen("tcp", s)
		if err != nil {
			log.Fatalf("Failed to listen on %s: %v", s, err)
		}
		defer listener.Close()
		log.Info("Listening on ", s)
	*/

	for {
		// Accept incoming connection
		src, err := gus.Accept()
		if err != nil {
			log.Printf("Failed to accept client connection: %v\n", err)
			return
		}

		// make a new event from this new connection
		log.Printf("/me Connected on port %d from client: %s to where = %s\n", localport, src.RemoteAddr(), where)
		fmt.Fprintln(src, "/me hello")

		reader := bufio.NewReader(src)

		// Read one line
		line, err := reader.ReadString('\n')
		if err != nil {
			log.Info("Error reading line:", err)
			return
		}
		log.Info("gus got Received:", line)
		parts := strings.Fields(line)
		if len(parts) != 3 {
			return
		}
		if parts[0] != "/me" {
			return
		}
		if parts[1] != "hostname" {
			return
		}
		hostname := parts[2]
		msg := fmt.Sprintf("got hostname %s for %s", hostname, src.RemoteAddr())
		log.Info("gus:", msg)
		fmt.Fprintln(src, msg)

		if hostname == "framebook.wit.com" {
			// Handle the connection in a separate goroutine
			go handleConnection(src, pm.Dest, int(pm.Localport))
		}
	}
}

func gus3000(pm *Portmap) {
	port := int(pm.Localport)
	connect := pm.Dest

	// Listen on local port 3000
	s := fmt.Sprintf("0.0.0.0:%d", port)
	listener, err := net.Listen("tcp", s)
	if err != nil {
		log.Fatalf("Failed to listen on %s: %v", s, err)
	}
	defer listener.Close()
	log.Info("Listening on ", s)

	if pm.UseME {
		doME(pm, listener)
		return
	}

	for {
		// Accept incoming connection
		clientConn, err := listener.Accept()
		if err != nil {
			log.Printf("Failed to accept client connection: %v", err)
			continue
		}
		// log.Printf("Client connected: %s", clientConn.RemoteAddr())

		// Handle the connection in a separate goroutine
		go handleConnection(clientConn, connect, port)
	}
}

func handleConnection(clientConn net.Conn, where string, localport int) {
	defer clientConn.Close()

	// Connect to the target server
	// targetConn, err := net.Dial("tcp", "go.wit.com:80")
	targetConn, err := net.Dial("tcp", where)
	if err != nil {
		log.Printf("Failed to connect to %s %v", where, err)
		return
	}
	defer targetConn.Close()

	// make a new event from this new connection
	log.Printf("Connected on port %d from client: %s to where = %s\n", localport, clientConn.RemoteAddr(), where)
	e := new(Event)
	e.Etype = GusEventType_Connect
	e.LocalPort = int64(localport)
	e.Sock = new(GusSocket)
	e.Sock.SrcIp = fmt.Sprintf("%s", clientConn.RemoteAddr())
	e.Sock.DestIp = where
	e.Ctime = timestamppb.New(time.Now())
	me.events.Append(e)
	me.eventsChanged = true

	// Bidirectional copy of data
	// go io.Copy(targetConn, clientConn) // Client -> Target
	// io.Copy(clientConn, targetConn)    // Target -> Client
	ioCopy(clientConn, targetConn)

	// if the socket closes, record the close time
	e.Etime = timestamppb.New(time.Now())
	me.eventsChanged = true
	log.Printf("Connection closed on port %d from client: %s to where = %s\n", localport, clientConn.RemoteAddr(), where)
}

func enablePort(port int, dest string) {
	var found bool
	all := me.portmaps.All()
	for all.Scan() {
		pm := all.Next()
		if int(pm.Localport) == port {
			found = true
			log.Info("Found port!", port)
			if pm.Enabled {
				log.Info("portmap already enabled for", pm.Localport, "to", pm.Dest)
			} else {
				log.Info("portmap not enabled for", pm.Localport, "to", pm.Dest)
			}
		}
		// go gus3000(pm)
	}
	if !found {
		log.Info("Did not find port =", port)
	}
}

// should cleanly close both
func ioCopy(clientConn, targetConn net.Conn) {
	defer clientConn.Close()
	defer targetConn.Close()

	var wg sync.WaitGroup
	wg.Add(2)

	// Copy data in both directions
	go func() {
		defer wg.Done()
		io.Copy(targetConn, clientConn) // Client -> Target
		targetConn.Close()              // Ensure closure on EOF/error
	}()

	go func() {
		defer wg.Done()
		io.Copy(clientConn, targetConn) // Target -> Client
		clientConn.Close()              // Ensure closure on EOF/error
	}()

	wg.Wait() // Wait for both copies to complete
}