main.go

package main

import (
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/gofrs/flock"
	"github.com/spf13/pflag"
)

var (
	wg       = sync.WaitGroup{}
	helpFlag = pflag.BoolP("help", "p", false, "show help")
)

func main() {
	socketPath := "/tmp/rsyslog.sock"
	pflag.Parse()
	if *helpFlag {
		fmt.Printf("%s: \nObtain SSHD logs from rsyslog, record them, and provide query services.\n", os.Args[0])
		fmt.Printf("You need to forward the sshd logs to the Unix socket(%s) via rsyslog.\n", socketPath)
		fmt.Printf("You can access the Unix socket(%s) to obtain online user information.\n", GIN_SOCK_PATH)
		fmt.Println("such as: curl -s --unix-socket /tmp//tmp/sshd-tool.sock http://localhost/user | jq")
		return
	}

	fileLock := flock.New("/var/lock/sshd-tool.lock", flock.SetPermissions(0644))
	l, err := fileLock.TryLock()
	if err != nil {
		log.Fatalf("error lock file /var/lock/sshd-tool.lock, %v", err)
	}
	if !l {
		log.Fatalf("can't lock /var/lock/sshd-tool.lock, Perhaps an instance is already running.")
	}

	InitSSH()
	serverListener := InitGin()

	err = os.RemoveAll(socketPath)

	if err != nil {
		log.Fatalf("error delete %s: %v", socketPath, err)
	}

	conn, err := net.ListenPacket("unixgram", socketPath)

	if err != nil {
		log.Fatalf("listen unix socket %s failed: %s", socketPath, err.Error())
	}
	err = os.Chmod(socketPath, 0666)
	if err != nil {
		log.Fatalf("chmod 666 %s failed: %v", socketPath, err)
	}

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

	wg.Add(1)
	go func(wg *sync.WaitGroup) {
		select {
		case <-sigChan:
			conn.Close()
			serverListener.Close()
			os.Remove(socketPath)
			os.Remove(GIN_SOCK_PATH)
			fileLock.Unlock()
		case <-globalCtx.Done():
			conn.Close()
			serverListener.Close()
			os.Remove(socketPath)
			os.Remove(GIN_SOCK_PATH)
			fileLock.Unlock()
		}
		wg.Done()
	}(&wg)

	buffer := make([]byte, 16384)
	for {
		n, _, err := conn.ReadFrom(buffer)
		if err != nil {
			globalCancelFunc()
			break
		}
		go ParseSSHLog(string(buffer[:n]))
	}
	wg.Wait()
}