+56
-3
cmd/atscanner.go
+56
-3
cmd/atscanner.go
···
3
3
import (
4
4
"context"
5
5
"flag"
6
+
"fmt"
6
7
"os"
7
8
"os/signal"
8
9
"syscall"
···
17
18
"github.com/atscan/atscanner/internal/worker"
18
19
)
19
20
21
+
const VERSION = "1.0.0"
22
+
20
23
func main() {
21
24
configPath := flag.String("config", "config.yaml", "path to config file")
22
25
verbose := flag.Bool("verbose", false, "enable verbose logging")
···
25
28
// Load configuration
26
29
cfg, err := config.Load(*configPath)
27
30
if err != nil {
28
-
log.Fatal("Failed to load config: %v", err)
31
+
fmt.Fprintf(os.Stderr, "Failed to load config: %v\n", err)
32
+
os.Exit(1)
29
33
}
30
34
31
35
// Override verbose setting if flag is provided
···
36
40
// Initialize logger
37
41
log.Init(cfg.API.Verbose)
38
42
43
+
// Print banner
44
+
log.Banner(VERSION)
45
+
46
+
// Print configuration summary
47
+
log.PrintConfig(map[string]string{
48
+
"Database Type": cfg.Database.Type,
49
+
"Database Path": cfg.Database.Path, // Will be auto-redacted
50
+
"PLC Directory": cfg.PLC.DirectoryURL,
51
+
"PLC Scan Interval": cfg.PLC.ScanInterval.String(),
52
+
"PLC Bundle Dir": cfg.PLC.BundleDir,
53
+
"PLC Cache": fmt.Sprintf("%v", cfg.PLC.UseCache),
54
+
"PLC Index DIDs": fmt.Sprintf("%v", cfg.PLC.IndexDIDs),
55
+
"PDS Scan Interval": cfg.PDS.ScanInterval.String(),
56
+
"PDS Workers": fmt.Sprintf("%d", cfg.PDS.Workers),
57
+
"PDS Timeout": cfg.PDS.Timeout.String(),
58
+
"API Host": cfg.API.Host,
59
+
"API Port": fmt.Sprintf("%d", cfg.API.Port),
60
+
"Verbose Logging": fmt.Sprintf("%v", cfg.API.Verbose),
61
+
})
62
+
39
63
// Initialize database using factory pattern
40
64
db, err := storage.NewDatabase(cfg.Database.Type, cfg.Database.Path)
41
65
if err != nil {
42
66
log.Fatal("Failed to initialize database: %v", err)
43
67
}
44
-
defer db.Close()
68
+
defer func() {
69
+
log.Info("Closing database connection...")
70
+
db.Close()
71
+
}()
45
72
46
73
// Set scan retention from config
47
-
db.SetScanRetention(cfg.PDS.ScanRetention)
74
+
if cfg.PDS.ScanRetention > 0 {
75
+
db.SetScanRetention(cfg.PDS.ScanRetention)
76
+
log.Verbose("Scan retention set to %d scans per endpoint", cfg.PDS.ScanRetention)
77
+
}
48
78
49
79
// Run migrations
50
80
if err := db.Migrate(); err != nil {
···
55
85
defer cancel()
56
86
57
87
// Initialize workers
88
+
log.Info("Initializing scanners...")
89
+
58
90
plcScanner := plc.NewScanner(db, cfg.PLC)
59
91
defer plcScanner.Close()
92
+
log.Verbose("✓ PLC scanner initialized")
60
93
61
94
pdsScanner := pds.NewScanner(db, cfg.PDS)
95
+
log.Verbose("✓ PDS scanner initialized")
62
96
63
97
scheduler := worker.NewScheduler()
64
98
···
68
102
log.Error("PLC scan error: %v", err)
69
103
}
70
104
})
105
+
log.Verbose("✓ PLC scan job scheduled (interval: %s)", cfg.PLC.ScanInterval)
71
106
72
107
// Schedule PDS availability checks
73
108
scheduler.AddJob("pds_scan", cfg.PDS.ScanInterval, func() {
···
75
110
log.Error("PDS scan error: %v", err)
76
111
}
77
112
})
113
+
log.Verbose("✓ PDS scan job scheduled (interval: %s)", cfg.PDS.ScanInterval)
78
114
79
115
// Start API server
116
+
log.Info("Starting API server on %s:%d...", cfg.API.Host, cfg.API.Port)
80
117
apiServer := api.NewServer(db, cfg.API, cfg.PLC)
81
118
go func() {
82
119
if err := apiServer.Start(); err != nil {
···
84
121
}
85
122
}()
86
123
124
+
// Give the API server a moment to start
125
+
time.Sleep(100 * time.Millisecond)
126
+
log.Info("✓ API server started successfully")
127
+
log.Info("")
128
+
log.Info("🚀 ATScanner is running!")
129
+
log.Info(" API available at: http://%s:%d", cfg.API.Host, cfg.API.Port)
130
+
log.Info(" Press Ctrl+C to stop")
131
+
log.Info("")
132
+
87
133
// Start scheduler
88
134
scheduler.Start(ctx)
89
135
···
92
138
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
93
139
<-sigChan
94
140
141
+
log.Info("")
95
142
log.Info("Shutting down gracefully...")
96
143
cancel()
144
+
145
+
log.Info("Stopping API server...")
97
146
apiServer.Shutdown(context.Background())
147
+
148
+
log.Info("Waiting for active tasks to complete...")
98
149
time.Sleep(2 * time.Second)
150
+
151
+
log.Info("✓ Shutdown complete. Goodbye!")
99
152
}
+97
internal/log/log.go
+97
internal/log/log.go
···
1
1
package log
2
2
3
3
import (
4
+
"fmt"
4
5
"io"
5
6
"log"
6
7
"os"
8
+
"strings"
7
9
)
8
10
9
11
var (
···
39
41
func Fatal(format string, v ...interface{}) {
40
42
errorLog.Fatalf(format, v...)
41
43
}
44
+
45
+
// Banner prints a startup banner
46
+
func Banner(version string) {
47
+
banner := `
48
+
╔════════════════════════════════════════════════════════════╗
49
+
║ ║
50
+
║ █████╗ ████████╗███████╗ ██████╗ █████╗ ███╗ ██╗ ║
51
+
║ ██╔══██╗╚══██╔══╝██╔════╝██╔════╝██╔══██╗████╗ ██║ ║
52
+
║ ███████║ ██║ ███████╗██║ ███████║██╔██╗ ██║ ║
53
+
║ ██╔══██║ ██║ ╚════██║██║ ██╔══██║██║╚██╗██║ ║
54
+
║ ██║ ██║ ██║ ███████║╚██████╗██║ ██║██║ ╚████║ ║
55
+
║ ╚═╝ ╚═╝ ╚═╝ ╚══════╝ ╚═════╝╚═╝ ╚═╝╚═╝ ╚═══╝ ║
56
+
║ ║
57
+
║ AT Protocol Network Scanner & Indexer ║
58
+
║ Version %s ║
59
+
║ ║
60
+
╚════════════════════════════════════════════════════════════╝
61
+
`
62
+
fmt.Printf(banner, padVersion(version))
63
+
}
64
+
65
+
// padVersion pads the version string to fit the banner
66
+
func padVersion(version string) string {
67
+
targetLen := 7
68
+
if len(version) < targetLen {
69
+
padding := strings.Repeat(" ", (targetLen-len(version))/2)
70
+
return padding + version + padding
71
+
}
72
+
return version
73
+
}
74
+
75
+
// RedactPassword redacts passwords from connection strings
76
+
func RedactPassword(connStr string) string {
77
+
// Handle PostgreSQL URI format: postgresql://user:password@host/db
78
+
// Pattern: find everything between :// and @ that contains a colon
79
+
if strings.Contains(connStr, "://") && strings.Contains(connStr, "@") {
80
+
// Find the credentials section
81
+
parts := strings.SplitN(connStr, "://", 2)
82
+
if len(parts) == 2 {
83
+
scheme := parts[0]
84
+
remainder := parts[1]
85
+
86
+
// Find the @ symbol
87
+
atIndex := strings.Index(remainder, "@")
88
+
if atIndex > 0 {
89
+
credentials := remainder[:atIndex]
90
+
hostAndDb := remainder[atIndex:]
91
+
92
+
// Check if there's a password (look for colon in credentials)
93
+
colonIndex := strings.Index(credentials, ":")
94
+
if colonIndex > 0 {
95
+
username := credentials[:colonIndex]
96
+
return fmt.Sprintf("%s://%s:***%s", scheme, username, hostAndDb)
97
+
}
98
+
}
99
+
}
100
+
}
101
+
102
+
// Handle key-value format: host=localhost password=secret user=myuser
103
+
if strings.Contains(connStr, "password=") {
104
+
parts := strings.Split(connStr, " ")
105
+
for i, part := range parts {
106
+
if strings.HasPrefix(part, "password=") {
107
+
parts[i] = "password=***"
108
+
}
109
+
}
110
+
return strings.Join(parts, " ")
111
+
}
112
+
113
+
return connStr
114
+
}
115
+
116
+
// PrintConfig prints configuration summary
117
+
func PrintConfig(items map[string]string) {
118
+
Info("=== Configuration ===")
119
+
maxKeyLen := 0
120
+
for key := range items {
121
+
if len(key) > maxKeyLen {
122
+
maxKeyLen = len(key)
123
+
}
124
+
}
125
+
126
+
for key, value := range items {
127
+
padding := strings.Repeat(" ", maxKeyLen-len(key))
128
+
129
+
// Redact database connection strings
130
+
displayValue := value
131
+
if strings.Contains(key, "Database Path") || strings.Contains(key, "Connection") || strings.Contains(strings.ToLower(key), "password") {
132
+
displayValue = RedactPassword(value)
133
+
}
134
+
135
+
fmt.Printf(" %s:%s %s\n", key, padding, displayValue)
136
+
}
137
+
Info("====================")
138
+
}
+21
-3
internal/storage/postgres.go
+21
-3
internal/storage/postgres.go
···
8
8
"strings"
9
9
"time"
10
10
11
+
"github.com/atscan/atscanner/internal/log"
11
12
"github.com/jackc/pgx/v5"
12
13
"github.com/jackc/pgx/v5/pgxpool"
13
14
_ "github.com/jackc/pgx/v5/stdlib"
···
20
21
}
21
22
22
23
func NewPostgresDB(connString string) (*PostgresDB, error) {
24
+
log.Info("Connecting to PostgreSQL database...")
25
+
23
26
// Open standard sql.DB (for compatibility)
24
27
db, err := sql.Open("pgx", connString)
25
28
if err != nil {
26
-
return nil, err
29
+
return nil, fmt.Errorf("failed to open database: %w", err)
27
30
}
28
31
29
32
// Connection pool settings
···
31
34
db.SetMaxIdleConns(25)
32
35
db.SetConnMaxLifetime(5 * time.Minute)
33
36
db.SetConnMaxIdleTime(2 * time.Minute)
37
+
38
+
log.Verbose(" Max open connections: 50")
39
+
log.Verbose(" Max idle connections: 25")
40
+
log.Verbose(" Connection max lifetime: 5m")
34
41
35
42
// Test connection
43
+
log.Info("Testing database connection...")
36
44
if err := db.Ping(); err != nil {
37
45
return nil, fmt.Errorf("failed to ping database: %w", err)
38
46
}
47
+
log.Info("✓ Database connection successful")
39
48
40
49
// Also create pgx pool for COPY operations
50
+
log.Verbose("Creating pgx connection pool...")
41
51
pool, err := pgxpool.New(context.Background(), connString)
42
52
if err != nil {
43
53
return nil, fmt.Errorf("failed to create pgx pool: %w", err)
44
54
}
55
+
log.Verbose("✓ Connection pool created")
45
56
46
57
return &PostgresDB{
47
58
db: db,
48
59
pool: pool,
49
-
scanRetention: 3,
60
+
scanRetention: 3, // Default
50
61
}, nil
51
62
}
52
63
···
58
69
}
59
70
60
71
func (p *PostgresDB) Migrate() error {
72
+
log.Info("Running database migrations...")
73
+
61
74
schema := `
62
75
-- Endpoints table (NO user_count, NO ip_info)
63
76
CREATE TABLE IF NOT EXISTS endpoints (
···
180
193
`
181
194
182
195
_, err := p.db.Exec(schema)
183
-
return err
196
+
if err != nil {
197
+
return err
198
+
}
199
+
200
+
log.Info("✓ Database migrations completed successfully")
201
+
return nil
184
202
}
185
203
186
204
// ===== ENDPOINT OPERATIONS =====