a

app/services.go

@@ -10,6 +10,7 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"path/filepath"
 	"strconv"
 	"strings"
 	"sync"
@@ -26,11 +27,19 @@ type App struct {
     cancel      context.CancelFunc
     wg          sync.WaitGroup
     // keep references to servers for graceful shutdown
-    httpSrv  *http.Server
+    httpSrvs  []*http.Server
     sshSigner ssh.Signer
     // track TCP listeners for graceful shutdown
     mu        sync.Mutex
     listeners []net.Listener
+    conns     map[net.Conn]struct{}
+}
+
+// addHTTPServer registers an HTTP server for graceful shutdown
+func (a *App) addHTTPServer(s *http.Server) {
+    a.mu.Lock()
+    a.httpSrvs = append(a.httpSrvs, s)
+    a.mu.Unlock()
 }
 
 // addListener registers a listener for later shutdown
@@ -123,32 +132,40 @@ func trimQuotes(s string) string {
 }
 
 func NewApp(cfg Config) (*App, error) {
-	l, err := NewLogger(cfg)
-	if err != nil {
-		return nil, err
-	}
-	ctx, cancel := context.WithCancel(context.Background())
-	a := &App{cfg: cfg, logger: l, ctx: ctx, cancel: cancel}
-	return a, nil
+    l, err := NewLogger(cfg)
+    if err != nil {
+        return nil, err
+    }
+    // Initialize threat intelligence system
+    threatIntelPath := "threat_intel.json"
+    if cfg.LogPath != "" {
+        threatIntelPath = filepath.Join(filepath.Dir(cfg.LogPath), "threat_intel.json")
+    }
+    ti := NewThreatIntel(threatIntelPath, true)
+
+    // Root context for the App used for shutdown signalling
+    ctx, cancel := context.WithCancel(context.Background())
+    a := &App{cfg: cfg, logger: l, threatIntel: ti, ctx: ctx, cancel: cancel, conns: make(map[net.Conn]struct{})}
+    return a, nil
 }
 
 func (a *App) Run(ctx context.Context) error {
-	// start services according to cfg
-	if a.cfg.Services.HTTP {
-		a.wg.Add(1)
-		go func() {
-			defer a.wg.Done()
-			a.startHTTP(a.cfg.Ports.HTTP)
-		}()
-	}
-	if a.cfg.Services.SSH {
-		a.wg.Add(1)
-		go func() {
-			defer a.wg.Done()
-			a.startTCPService("ssh", a.cfg.Ports.SSH, a.sshHandler)
-		}()
-	}
-	if a.cfg.Services.FTP {
+    // start services according to cfg
+    if a.cfg.Services.HTTP {
+        a.wg.Add(1)
+        go func() {
+            defer a.wg.Done()
+            a.startHTTP(a.cfg.Ports.HTTP)
+        }()
+    }
+    if a.cfg.Services.SSH {
+        a.wg.Add(1)
+        go func() {
+            defer a.wg.Done()
+            a.startTCPService("ssh", a.cfg.Ports.SSH, a.sshHandler)
+        }()
+    }
+    if a.cfg.Services.FTP {
 		a.wg.Add(1)
 		go func() {
 			defer a.wg.Done()
@@ -285,15 +302,27 @@ func (a *App) Run(ctx context.Context) error {
 
 func (a *App) Shutdown() {
     a.cancel()
-    // attempt to close http server if running
-    if a.httpSrv != nil {
+    // attempt to close all http servers if running
+    a.mu.Lock()
+    srvs := a.httpSrvs
+    a.httpSrvs = nil
+    a.mu.Unlock()
+    for _, s := range srvs {
         ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-        defer cancel()
-        _ = a.httpSrv.Shutdown(ctx)
+        _ = s.Shutdown(ctx)
+        cancel()
     }
     // close all TCP listeners to unblock Accept loops
     a.closeAllListeners()
+    // close all tracked connections to unblock handlers
+    a.closeAllConns()
     a.wg.Wait()
+    
+    // Save threat intelligence data before shutdown
+    if a.threatIntel != nil {
+        _ = a.threatIntel.Save()
+    }
+    
     _ = a.logger.Close()
 }
 
@@ -308,12 +337,47 @@ func (a *App) closeAllListeners() {
     }
 }
 
+// addConn tracks an active connection for shutdown
+func (a *App) addConn(c net.Conn) {
+    a.mu.Lock()
+    if a.conns == nil {
+        a.conns = make(map[net.Conn]struct{})
+    }
+    a.conns[c] = struct{}{}
+    a.mu.Unlock()
+}
+
+// removeConn removes a connection from tracking
+func (a *App) removeConn(c net.Conn) {
+    a.mu.Lock()
+    delete(a.conns, c)
+    a.mu.Unlock()
+}
+
+// closeAllConns closes all tracked connections to unblock handlers
+func (a *App) closeAllConns() {
+    a.mu.Lock()
+    conns := make([]net.Conn, 0, len(a.conns))
+    for c := range a.conns {
+        conns = append(conns, c)
+    }
+    a.conns = make(map[net.Conn]struct{})
+    a.mu.Unlock()
+    for _, c := range conns {
+        _ = c.Close()
+    }
+}
+
 // helpers for logging
 func (a *App) logEvent(r Record) {
-    if a.logger == nil {
-        return
+    if a.logger != nil {
+        _ = a.logger.Log(r)
+    }
+    
+    // Record threat intelligence if IP is not private
+    if a.threatIntel != nil && r.RemoteAddr != "" && !IsPrivateIP(r.RemoteAddr) {
+        a.threatIntel.RecordActivity(r)
     }
-    _ = a.logger.Log(r)
 }
 
 // HTTP honeypot
@@ -344,7 +408,7 @@ func (a *App) startHTTP(port int) {
     })
 
     srv := &http.Server{Addr: addr, Handler: mux}
-    a.httpSrv = srv
+    a.addHTTPServer(srv)
     log.Printf("HTTP listening on %s", addr)
     if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
         log.Printf("http server error: %v", err)
@@ -399,7 +463,12 @@ func (a *App) startTCPService(name string, port int, handler func(net.Conn)) {
             return
         case conn := <-acceptCh:
             go func(c net.Conn) {
-                defer c.Close()
+                // track and ensure cleanup
+                a.addConn(c)
+                defer func() {
+                    a.removeConn(c)
+                    _ = c.Close()
+                }()
                 // Check if context is cancelled before processing
                 select {
                 case <-a.ctx.Done():