-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathdb.go
More file actions
84 lines (71 loc) · 2.04 KB
/
db.go
File metadata and controls
84 lines (71 loc) · 2.04 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
package main
import (
"database/sql"
"fmt"
"log"
"time"
_ "modernc.org/sqlite" // Pure Go SQLite driver
)
var DB *sql.DB
// InitDB initializes the SQLite database and creates tables
func InitDB() {
var err error
// Open the database file (c2.db will be created automatically)
DB, err = sql.Open("sqlite", "c2.db")
if err != nil {
log.Fatal("[-] Database Connection Error: ", err)
}
// Create Agents Table
createTableSQL := `CREATE TABLE IF NOT EXISTS agents (
id TEXT PRIMARY KEY,
ip TEXT,
hostname TEXT,
platform TEXT,
username TEXT,
status TEXT,
last_seen DATETIME
);`
_, err = DB.Exec(createTableSQL)
if err != nil {
log.Fatal("[-] Error Creating Table: ", err)
}
fmt.Println("[+] Database initialized (c2.db)")
}
// UpsertAgent inserts a new agent or updates an existing one
func UpsertAgent(agent *Agent) {
// Use INSERT OR REPLACE logic for efficiency
query := `INSERT INTO agents (id, ip, hostname, platform, username, status, last_seen)
VALUES (?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
ip=excluded.ip,
hostname=excluded.hostname,
username=excluded.username,
status=excluded.status,
last_seen=excluded.last_seen;`
_, err := DB.Exec(query, agent.ID, agent.IP, agent.Hostname, agent.Platform, agent.Username, agent.Status, agent.LastSeen)
if err != nil {
log.Printf("[-] DB Error (UpsertAgent): %v\n", err)
}
}
// LoadAgents loads all agents from the database into memory on startup
func LoadAgents() map[string]*Agent {
loadedAgents := make(map[string]*Agent)
rows, err := DB.Query("SELECT id, ip, hostname, platform, username, status, last_seen FROM agents")
if err != nil {
log.Printf("[-] DB Error (Query): %v\n", err)
return loadedAgents
}
defer rows.Close()
for rows.Next() {
var a Agent
var lastSeen time.Time
err := rows.Scan(&a.ID, &a.IP, &a.Hostname, &a.Platform, &a.Username, &a.Status, &lastSeen)
if err != nil {
continue
}
a.LastSeen = lastSeen
a.CommandQ = []string{} // Initialize empty queue
loadedAgents[a.ID] = &a
}
return loadedAgents
}