pkg/server/state/sqlite_ledger.go · contraband

        ● Notice //GitHub (@shehzade) is temporarily suspended pending an automated false-positive review. All tools are hosted on this Code Mirror and available on request (contact); GitHub links will return once reinstated.
      

Multi-protocol network egress assessment and DLP verification framework featuring automated L4 sweeps, corporate proxy traversal, and covert protocol smuggling across 20+ channels.

$ git clone https://shehzade.io/repos/contraband.git Copy

Source ZIP (1.93 MB)

Code 174 Commits 1

← Back to directory

810 lines · 22.8 KB

Raw

1	package state
2
3	import (
4	"crypto/sha256"
5	"database/sql"
6	"encoding/hex"
7	"encoding/json"
8	"fmt"
9	"strings"
10	"sync"
11	"time"
12
13	"github.com/shehzade/contraband/pkg/common"
14	_ "modernc.org/sqlite" // Pure-Go SQLite driver for CGO_ENABLED=0 static compilation
15	)
16
17	// SQLiteLedger implements an ACID-compliant, persistent session ledger using pure-Go SQLite.
18	type SQLiteLedger struct {
19	mu sync.RWMutex
20	db *sql.DB
21	observer LedgerObserver
22	closed bool
23	}
24
25	// NewSQLiteLedger opens or creates a SQLite database at the specified path.
26	func NewSQLiteLedger(dbPath string) (*SQLiteLedger, error) {
27	db, err := sql.Open("sqlite", dbPath)
28	if err != nil {
29	return nil, fmt.Errorf("failed to open sqlite database at %s: %w", dbPath, err)
30	}
31
32	db.SetMaxOpenConns(1)
33
34	ledger := &SQLiteLedger{
35	db: db,
36	}
37
38	if err := ledger.initSchema(); err != nil {
39	_ = db.Close()
40	return nil, fmt.Errorf("failed to initialize sqlite schema: %w", err)
41	}
42
43	return ledger, nil
44	}
45
46	// SetObserver attaches an observer to receive real-time chunk and session events.
47	func (s *SQLiteLedger) SetObserver(obs LedgerObserver) {
48	s.mu.Lock()
49	defer s.mu.Unlock()
50	s.observer = obs
51	}
52
53	func (s *SQLiteLedger) initSchema() error {
54	schema := `
55	CREATE TABLE IF NOT EXISTS probes (
56	id TEXT PRIMARY KEY,
57	timestamp TEXT NOT NULL,
58	stage TEXT NOT NULL,
59	protocol TEXT NOT NULL,
60	target_port INTEGER NOT NULL,
61	target_host TEXT NOT NULL,
62	local_addr TEXT,
63	proxy_hop TEXT,
64	tls_profile TEXT,
65	ja3_string TEXT,
66	ja4_string TEXT,
67	payload_type TEXT,
68	carrier_type TEXT,
69	target_volume_bytes INTEGER,
70	bytes_sent INTEGER,
71	bytes_recv INTEGER,
72	throughput_kbps REAL,
73	payload_hash TEXT,
74	wire_hash TEXT,
75	verdict TEXT NOT NULL,
76	captured_hash TEXT,
77	error_message TEXT,
78	raw_evidence TEXT,
79	timestamps_json TEXT
80	);
81
82	CREATE TABLE IF NOT EXISTS sessions (
83	session_id TEXT PRIMARY KEY,
84	token_hex TEXT NOT NULL,
85	token TEXT,
86	client_ip TEXT NOT NULL,
87	protocol TEXT NOT NULL,
88	port INTEGER NOT NULL,
89	start_time TEXT NOT NULL,
90	last_active TEXT NOT NULL,
91	total_bytes INTEGER NOT NULL,
92	expected_bytes INTEGER NOT NULL,
93	total_chunks INTEGER NOT NULL,
94	received_chunks INTEGER NOT NULL,
95	payload_hash TEXT,
96	expected_hash TEXT,
97	status TEXT NOT NULL,
98	verdict TEXT NOT NULL
99	);
100
101	CREATE TABLE IF NOT EXISTS session_chunks (
102	session_id TEXT NOT NULL,
103	seq INTEGER NOT NULL,
104	total INTEGER NOT NULL,
105	data BLOB NOT NULL,
106	received_at TEXT NOT NULL,
107	PRIMARY KEY (session_id, seq),
108	FOREIGN KEY (session_id) REFERENCES sessions(session_id) ON DELETE CASCADE
109	);
110
111	CREATE VIEW IF NOT EXISTS chunks AS
112	SELECT session_id AS token, seq AS sequence, total, LENGTH(data) AS length, received_at, data AS payload
113	FROM session_chunks;
114
115	CREATE TRIGGER IF NOT EXISTS insert_chunks_trigger INSTEAD OF INSERT ON chunks
116	BEGIN
117	INSERT INTO session_chunks (session_id, seq, total, data, received_at)
118	VALUES (NEW.token, NEW.sequence, NEW.total, NEW.payload, NEW.received_at)
119	ON CONFLICT(session_id, seq) DO NOTHING;
120	END;
121
122	CREATE TABLE IF NOT EXISTS credentials (
123	id INTEGER PRIMARY KEY AUTOINCREMENT,
124	protocol TEXT NOT NULL,
125	timestamp TEXT NOT NULL,
126	client_ip TEXT NOT NULL,
127	username TEXT NOT NULL,
128	domain TEXT,
129	hash_type TEXT NOT NULL,
130	hash TEXT NOT NULL
131	);
132
133	CREATE INDEX IF NOT EXISTS idx_session_chunks_session ON session_chunks(session_id);
134	CREATE INDEX IF NOT EXISTS idx_probes_stage ON probes(stage);
135	CREATE INDEX IF NOT EXISTS idx_probes_verdict ON probes(verdict);
136	CREATE INDEX IF NOT EXISTS idx_probes_timestamp ON probes(timestamp);
137	`
138	_, err := s.db.Exec(schema)
139	return err
140	}
141
142	func (s *SQLiteLedger) RecordProbe(result common.ProbeResult) error {
143	s.mu.Lock()
144	defer s.mu.Unlock()
145
146	if s.closed {
147	return ErrLedgerClosed
148	}
149
150	if result.Timestamp.IsZero() {
151	result.Timestamp = time.Now().UTC()
152	}
153	if result.TargetPort == 0 {
154	result.TargetPort = 80
155	}
156
157	if err := result.Validate(); err != nil {
158	return fmt.Errorf("invalid probe result: %w", err)
159	}
160
161	tsJSON, err := json.Marshal(result.Timestamps)
162	if err != nil {
163	return fmt.Errorf("failed to marshal timestamps: %w", err)
164	}
165
166	query := `
167	INSERT INTO probes (
168	id, timestamp, stage, protocol, target_port, target_host,
169	local_addr, proxy_hop, tls_profile, ja3_string, ja4_string,
170	payload_type, carrier_type, target_volume_bytes, bytes_sent, bytes_recv,
171	throughput_kbps, payload_hash, wire_hash, verdict,
172	captured_hash, error_message, raw_evidence, timestamps_json
173	) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
174	ON CONFLICT(id) DO UPDATE SET
175	verdict = excluded.verdict,
176	bytes_sent = excluded.bytes_sent,
177	bytes_recv = excluded.bytes_recv,
178	throughput_kbps = excluded.throughput_kbps,
179	payload_hash = excluded.payload_hash,
180	wire_hash = excluded.wire_hash,
181	error_message = excluded.error_message,
182	raw_evidence = excluded.raw_evidence,
183	timestamps_json = excluded.timestamps_json;
184	`
185	_, err = s.db.Exec(query,
186	result.ID, result.Timestamp.Format(time.RFC3339Nano), result.Stage, result.Protocol,
187	result.TargetPort, result.TargetHost, result.LocalAddr, result.ProxyHop,
188	result.TLSProfile, result.JA3String, result.JA4String, result.PayloadType,
189	result.CarrierType, result.TargetVolume, result.BytesSent, result.BytesRecv,
190	result.ThroughputBps, result.PayloadHash, result.WireHash, string(result.Verdict),
191	result.CapturedHash, result.ErrorMessage, result.RawEvidence, string(tsJSON),
192	)
193	if err == nil && s.observer != nil {
194	if po, ok := s.observer.(ProbeObserver); ok {
195	po.OnProbeRecorded(&result)
196	}
197	}
198	return err
199	}
200
201	func (s SQLiteLedger) GetProbe(id string) (common.ProbeResult, error) {
202	s.mu.RLock()
203	defer s.mu.RUnlock()
204
205	if s.closed {
206	return nil, ErrLedgerClosed
207	}
208
209	query := `SELECT id, timestamp, stage, protocol, target_port, target_host, local_addr,
210	proxy_hop, tls_profile, ja3_string, ja4_string, payload_type, carrier_type,
211	target_volume_bytes, bytes_sent, bytes_recv, throughput_kbps, payload_hash,
212	wire_hash, verdict, captured_hash, error_message, raw_evidence, timestamps_json
213	FROM probes WHERE id = ?`
214
215	row := s.db.QueryRow(query, id)
216	return scanProbe(row)
217	}
218
219	func (s *SQLiteLedger) GetAllProbes() ([]common.ProbeResult, error) {
220	s.mu.RLock()
221	defer s.mu.RUnlock()
222
223	if s.closed {
224	return nil, ErrLedgerClosed
225	}
226
227	query := `SELECT id, timestamp, stage, protocol, target_port, target_host, local_addr,
228	proxy_hop, tls_profile, ja3_string, ja4_string, payload_type, carrier_type,
229	target_volume_bytes, bytes_sent, bytes_recv, throughput_kbps, payload_hash,
230	wire_hash, verdict, captured_hash, error_message, raw_evidence, timestamps_json
231	FROM probes ORDER BY timestamp ASC`
232
233	rows, err := s.db.Query(query)
234	if err != nil {
235	return nil, err
236	}
237	defer rows.Close()
238
239	var results []common.ProbeResult
240	for rows.Next() {
241	p, err := scanProbe(rows)
242	if err != nil {
243	return nil, err
244	}
245	results = append(results, *p)
246	}
247	return results, rows.Err()
248	}
249
250	func (s SQLiteLedger) GetSession(token common.SessionToken) (common.Session, error) {
251	s.mu.RLock()
252	defer s.mu.RUnlock()
253
254	if s.closed {
255	return nil, ErrLedgerClosed
256	}
257
258	tokStr := token.String()
259	query := `SELECT session_id, token_hex, token, client_ip, protocol, port, start_time, last_active,
260	total_bytes, expected_bytes, total_chunks, received_chunks,
261	payload_hash, expected_hash, status, verdict
262	FROM sessions WHERE session_id = ? OR token_hex = ? OR token = ?`
263
264	row := s.db.QueryRow(query, tokStr, tokStr, tokStr)
265	session, err := scanSession(row, token)
266	if err != nil {
267	return nil, err
268	}
269
270	chunkQuery := `SELECT seq, total, LENGTH(data), received_at FROM session_chunks WHERE session_id = ?`
271	cRows, err := s.db.Query(chunkQuery, session.SessionID)
272	if err == nil {
273	defer cRows.Close()
274	for cRows.Next() {
275	var cr common.ChunkRecord
276	var rAt string
277	if err := cRows.Scan(&cr.Sequence, &cr.Total, &cr.Length, &rAt); err == nil {
278	cr.ReceivedAt, _ = time.Parse(time.RFC3339Nano, rAt)
279	session.Chunks[cr.Sequence] = cr
280	}
281	}
282	}
283
284	return session, nil
285	}
286
287	func (s SQLiteLedger) UpdateSession(session common.Session) error {
288	s.mu.Lock()
289	defer s.mu.Unlock()
290
291	if s.closed {
292	return ErrLedgerClosed
293	}
294	if session == nil {
295	return fmt.Errorf("cannot update nil session")
296	}
297
298	sessID := session.SessionID
299	tokStr := session.Token.String()
300	if sessID == "" {
301	sessID = tokStr
302	}
303	if tokStr == "" {
304	tokStr = sessID
305	}
306
307	query := `INSERT INTO sessions (
308	session_id, token_hex, token, client_ip, protocol, port,
309	start_time, last_active, total_bytes, expected_bytes,
310	total_chunks, received_chunks, payload_hash, expected_hash,
311	status, verdict
312	) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
313	ON CONFLICT(session_id) DO UPDATE SET
314	status = excluded.status,
315	verdict = excluded.verdict,
316	last_active = excluded.last_active,
317	payload_hash = CASE WHEN excluded.payload_hash != '' THEN excluded.payload_hash ELSE sessions.payload_hash END,
318	expected_hash = CASE WHEN excluded.expected_hash != '' THEN excluded.expected_hash ELSE sessions.expected_hash END,
319	total_bytes = CASE WHEN excluded.total_bytes > sessions.total_bytes THEN excluded.total_bytes ELSE sessions.total_bytes END,
320	received_chunks = CASE WHEN excluded.received_chunks > sessions.received_chunks THEN excluded.received_chunks ELSE sessions.received_chunks END,
321	total_chunks = CASE WHEN excluded.total_chunks > 0 THEN excluded.total_chunks ELSE sessions.total_chunks END;`
322
323	_, err := s.db.Exec(query,
324	sessID, tokStr, tokStr, session.ClientIP, session.Protocol, session.Port,
325	session.StartTime.Format(time.RFC3339Nano),
326	session.LastActive.Format(time.RFC3339Nano),
327	session.TotalBytes, session.ExpectedBytes,
328	session.TotalChunks, session.ReceivedChunks,
329	session.PayloadHash, session.ExpectedHash,
330	session.Status, string(session.Verdict),
331	)
332	return err
333	}
334
335	func (s SQLiteLedger) GetAllSessions() ([]common.Session, error) {
336	s.mu.RLock()
337	defer s.mu.RUnlock()
338
339	if s.closed {
340	return nil, ErrLedgerClosed
341	}
342
343	query := `SELECT session_id, token_hex, token, client_ip, protocol, port, start_time, last_active,
344	total_bytes, expected_bytes, total_chunks, received_chunks,
345	payload_hash, expected_hash, status, verdict
346	FROM sessions ORDER BY start_time ASC`
347
348	rows, err := s.db.Query(query)
349	if err != nil {
350	return nil, err
351	}
352	defer rows.Close()
353
354	var sessions []*common.Session
355	for rows.Next() {
356	sess, err := scanSession(rows, common.SessionToken{})
357	if err != nil {
358	return nil, err
359	}
360	sessions = append(sessions, sess)
361	}
362	return sessions, rows.Err()
363	}
364
365	func (s SQLiteLedger) AssembleChunk(token common.SessionToken, seq uint32, total uint32, payload []byte, expectedHash string) (common.Session, bool, error) {
366	return s.AssembleChunkWithMeta(token, seq, total, payload, expectedHash, "", "", 0, "", "")
367	}
368
369	func (s SQLiteLedger) AssembleChunkWithMeta(token common.SessionToken, seq uint32, total uint32, payload []byte, expectedHash, clientIP, protocol string, port int, clientHost, networkZone string) (common.Session, bool, error) {
370	s.mu.Lock()
371
372	if s.closed {
373	s.mu.Unlock()
374	return nil, false, ErrLedgerClosed
375	}
376	if total == 0 \|\| seq == 0 \|\| seq > total {
377	s.mu.Unlock()
378	return nil, false, fmt.Errorf("%w: seq=%d total=%d", ErrInvalidChunk, seq, total)
379	}
380
381	tokStr := token.String()
382	now := time.Now().UTC()
383	nowStr := now.Format(time.RFC3339Nano)
384
385	tx, err := s.db.Begin()
386	if err != nil {
387	s.mu.Unlock()
388	return nil, false, err
389	}
390	defer func() { _ = tx.Rollback() }()
391
392	upsertSession := `
393	INSERT INTO sessions (
394	session_id, token_hex, token, client_ip, protocol, port, start_time, last_active,
395	total_bytes, expected_bytes, total_chunks, received_chunks,
396	payload_hash, expected_hash, status, verdict
397	) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0, 0, ?, 0, '', ?, 'active', 'BLOCKED_L7')
398	ON CONFLICT(session_id) DO UPDATE SET
399	total_chunks = excluded.total_chunks,
400	client_ip = CASE WHEN sessions.client_ip = '' AND excluded.client_ip != '' THEN excluded.client_ip ELSE sessions.client_ip END,
401	protocol = CASE WHEN sessions.protocol = '' AND excluded.protocol != '' THEN excluded.protocol ELSE sessions.protocol END,
402	port = CASE WHEN sessions.port = 0 AND excluded.port != 0 THEN excluded.port ELSE sessions.port END,
403	expected_hash = CASE WHEN excluded.expected_hash != '' THEN excluded.expected_hash ELSE sessions.expected_hash END;
404	`
405	if _, err := tx.Exec(upsertSession, tokStr, tokStr, tokStr, clientIP, protocol, port, nowStr, nowStr, total, expectedHash); err != nil {
406	s.mu.Unlock()
407	return nil, false, err
408	}
409
410	insertChunk := `
411	INSERT INTO session_chunks (session_id, seq, total, data, received_at)
412	VALUES (?, ?, ?, ?, ?)
413	ON CONFLICT(session_id, seq) DO NOTHING;
414	`
415	res, err := tx.Exec(insertChunk, tokStr, seq, total, payload, nowStr)
416	if err != nil {
417	s.mu.Unlock()
418	return nil, false, err
419	}
420
421	rowsAffected, _ := res.RowsAffected()
422	if rowsAffected > 0 {
423	updateStats := `
424	UPDATE sessions SET
425	received_chunks = received_chunks + 1,
426	total_bytes = total_bytes + ?,
427	last_active = ?
428	WHERE session_id = ?;
429	`
430	if _, err := tx.Exec(updateStats, len(payload), nowStr, tokStr); err != nil {
431	s.mu.Unlock()
432	return nil, false, err
433	}
434	}
435
436	var sess common.Session
437	var sessID, tokHex string
438	var tokNull sql.NullString
439	var sTime, lActive, verdStr string
440	querySess := `SELECT session_id, token_hex, token, client_ip, protocol, port, start_time, last_active,
441	total_bytes, expected_bytes, total_chunks, received_chunks,
442	payload_hash, expected_hash, status, verdict
443	FROM sessions WHERE session_id = ?`
444	err = tx.QueryRow(querySess, tokStr).Scan(
445	&sessID, &tokHex, &tokNull, &sess.ClientIP, &sess.Protocol, &sess.Port, &sTime, &lActive,
446	&sess.TotalBytes, &sess.ExpectedBytes, &sess.TotalChunks, &sess.ReceivedChunks,
447	&sess.PayloadHash, &sess.ExpectedHash, &sess.Status, &verdStr,
448	)
449	if err != nil {
450	s.mu.Unlock()
451	return nil, false, err
452	}
453	sess.Token = token
454	sess.SessionID = sessID
455	sess.Verdict = common.Verdict(verdStr)
456	sess.StartTime, _ = time.Parse(time.RFC3339Nano, sTime)
457	sess.LastActive, _ = time.Parse(time.RFC3339Nano, lActive)
458
459	obs := s.observer
460
461	// Idempotency: if session was already completed or corrupted, return existing state
462	if sess.Status == "completed" {
463	_ = tx.Rollback()
464	retSess := CloneSession(&sess)
465	s.mu.Unlock()
466	return retSess, true, nil
467	}
468	if sess.Status == "corrupted" {
469	_ = tx.Rollback()
470	retSess := CloneSession(&sess)
471	s.mu.Unlock()
472	return retSess, true, ErrHashMismatch
473	}
474
475	if sess.ReceivedChunks >= int(total) {
476	chunkQuery := `SELECT seq, data FROM session_chunks WHERE session_id = ? ORDER BY seq ASC`
477	cRows, err := tx.Query(chunkQuery, tokStr)
478	if err != nil {
479	s.mu.Unlock()
480	return nil, false, err
481	}
482	defer cRows.Close()
483
484	chunkMap := make(map[uint32][]byte)
485	for cRows.Next() {
486	var sNum uint32
487	var pBytes []byte
488	if err := cRows.Scan(&sNum, &pBytes); err == nil {
489	chunkMap[sNum] = pBytes
490	}
491	}
492
493	res := ReassembleChunks(chunkMap, total, sess.ExpectedHash)
494	sess.PayloadHash = res.ComputedHash
495	sess.Status = res.Status
496	sess.Verdict = res.Verdict
497
498	updateFinal := `UPDATE sessions SET payload_hash = ?, status = ?, verdict = ? WHERE session_id = ?`
499	if _, err := tx.Exec(updateFinal, sess.PayloadHash, sess.Status, string(sess.Verdict), tokStr); err != nil {
500	s.mu.Unlock()
501	return nil, false, err
502	}
503
504	if err := tx.Commit(); err != nil {
505	s.mu.Unlock()
506	return nil, false, err
507	}
508
509	retSess := CloneSession(&sess)
510	s.mu.Unlock()
511
512	if obs != nil {
513	obs.OnChunkReceived(retSess, seq, total, int64(len(payload)))
514	obs.OnSessionCompleted(retSess)
515	}
516
517	if res.Error != nil {
518	return retSess, true, res.Error
519	}
520	return retSess, true, nil
521	}
522
523	if err := tx.Commit(); err != nil {
524	s.mu.Unlock()
525	return nil, false, err
526	}
527
528	retSess := CloneSession(&sess)
529	s.mu.Unlock()
530
531	if obs != nil {
532	obs.OnChunkReceived(retSess, seq, total, int64(len(payload)))
533	}
534
535	return retSess, false, nil
536	}
537
538	func (s *SQLiteLedger) GetSessionPayload(token common.SessionToken) ([]byte, error) {
539	s.mu.RLock()
540	defer s.mu.RUnlock()
541
542	if s.closed {
543	return nil, ErrLedgerClosed
544	}
545
546	tokStr := token.String()
547
548	var totalChunks int
549	var status, verdict, expectedHash string
550	err := s.db.QueryRow("SELECT total_chunks, status, verdict, expected_hash FROM sessions WHERE session_id = ? OR token_hex = ? OR token = ?", tokStr, tokStr, tokStr).Scan(&totalChunks, &status, &verdict, &expectedHash)
551	if err != nil {
552	if err == sql.ErrNoRows {
553	return nil, ErrSessionNotFound
554	}
555	return nil, err
556	}
557
558	if status == "corrupted" \|\| common.Verdict(verdict) == common.VerdictBlockedDLP {
559	return nil, ErrHashMismatch
560	}
561
562	if totalChunks <= 0 {
563	return nil, fmt.Errorf("%w: total chunks is %d", ErrSessionIncomplete, totalChunks)
564	}
565
566	var count int
567	if err := s.db.QueryRow("SELECT COUNT(*) FROM session_chunks WHERE session_id = ? OR session_id IN (SELECT session_id FROM sessions WHERE token_hex = ? OR token = ?)", tokStr, tokStr, tokStr).Scan(&count); err != nil {
568	return nil, err
569	}
570	if count < totalChunks {
571	return nil, fmt.Errorf("%w: received %d of %d chunks", ErrSessionIncomplete, count, totalChunks)
572	}
573
574	query := `SELECT seq, data FROM session_chunks WHERE session_id = ? OR session_id IN (SELECT session_id FROM sessions WHERE token_hex = ? OR token = ?) ORDER BY seq ASC`
575	rows, err := s.db.Query(query, tokStr, tokStr, tokStr)
576	if err != nil {
577	return nil, err
578	}
579	defer rows.Close()
580
581	var fullPayload []byte
582	var expectedSeq uint32 = 1
583
584	for rows.Next() {
585	var seq uint32
586	var payload []byte
587	if err := rows.Scan(&seq, &payload); err != nil {
588	return nil, err
589	}
590	if seq != expectedSeq {
591	return nil, fmt.Errorf("%w: missing sequence %d (found %d)", ErrSessionIncomplete, expectedSeq, seq)
592	}
593	fullPayload = append(fullPayload, payload...)
594	expectedSeq++
595	}
596
597	if int(expectedSeq-1) != totalChunks {
598	return nil, fmt.Errorf("%w: missing trailing chunks: expected %d, assembled %d", ErrSessionIncomplete, totalChunks, expectedSeq-1)
599	}
600
601	if expectedHash != "" {
602	hasher := sha256.New()
603	hasher.Write(fullPayload)
604	computed := hex.EncodeToString(hasher.Sum(nil))
605	if !strings.EqualFold(computed, expectedHash) {
606	return nil, ErrHashMismatch
607	}
608	}
609
610	return fullPayload, nil
611	}
612
613	func (s *SQLiteLedger) RecordCredential(cred common.CapturedCredential) error {
614	s.mu.Lock()
615	defer s.mu.Unlock()
616
617	if s.closed {
618	return ErrLedgerClosed
619	}
620
621	query := `INSERT INTO credentials (protocol, timestamp, client_ip, username, domain, hash_type, hash)
622	VALUES (?, ?, ?, ?, ?, ?, ?)`
623
624	_, err := s.db.Exec(query,
625	cred.Protocol, cred.Timestamp.Format(time.RFC3339Nano), cred.ClientIP,
626	cred.Username, cred.Domain, cred.HashType, cred.Hash,
627	)
628	return err
629	}
630
631	func (s *SQLiteLedger) GetCredentials() ([]common.CapturedCredential, error) {
632	s.mu.RLock()
633	defer s.mu.RUnlock()
634
635	if s.closed {
636	return nil, ErrLedgerClosed
637	}
638
639	query := `SELECT protocol, timestamp, client_ip, username, domain, hash_type, hash FROM credentials ORDER BY id ASC`
640	rows, err := s.db.Query(query)
641	if err != nil {
642	return nil, err
643	}
644	defer rows.Close()
645
646	var creds []common.CapturedCredential
647	for rows.Next() {
648	var c common.CapturedCredential
649	var ts string
650	if err := rows.Scan(&c.Protocol, &ts, &c.ClientIP, &c.Username, &c.Domain, &c.HashType, &c.Hash); err != nil {
651	return nil, err
652	}
653	c.Timestamp, _ = time.Parse(time.RFC3339Nano, ts)
654	creds = append(creds, c)
655	}
656	return creds, rows.Err()
657	}
658
659	func (s *SQLiteLedger) ExportJSON() ([]byte, error) {
660	probes, err := s.GetAllProbes()
661	if err != nil {
662	return nil, err
663	}
664	sessions, err := s.GetAllSessions()
665	if err != nil {
666	return nil, err
667	}
668	creds, err := s.GetCredentials()
669	if err != nil {
670	return nil, err
671	}
672
673	dump := struct {
674	TotalProbes int `json:"total_probes"`
675	TotalSessions int `json:"total_sessions"`
676	TotalCredentials int `json:"total_credentials"`
677	Probes []common.ProbeResult `json:"probes"`
678	Sessions []*common.Session `json:"sessions"`
679	Credentials []common.CapturedCredential `json:"credentials"`
680	}{
681	TotalProbes: len(probes),
682	TotalSessions: len(sessions),
683	TotalCredentials: len(creds),
684	Probes: probes,
685	Sessions: sessions,
686	Credentials: creds,
687	}
688
689	return json.MarshalIndent(dump, "", " ")
690	}
691
692	// Reset clears all records from session_chunks, sessions, probes, and credentials tables.
693	func (s *SQLiteLedger) Reset() error {
694	s.mu.Lock()
695	defer s.mu.Unlock()
696
697	if s.closed {
698	return ErrLedgerClosed
699	}
700
701	tx, err := s.db.Begin()
702	if err != nil {
703	return err
704	}
705	defer tx.Rollback()
706
707	if _, err := tx.Exec(`DELETE FROM session_chunks`); err != nil {
708	return err
709	}
710	if _, err := tx.Exec(`DELETE FROM sessions`); err != nil {
711	return err
712	}
713	if _, err := tx.Exec(`DELETE FROM probes`); err != nil {
714	return err
715	}
716	if _, err := tx.Exec(`DELETE FROM credentials`); err != nil {
717	return err
718	}
719
720	return tx.Commit()
721	}
722
723	func (s *SQLiteLedger) Close() error {
724	s.mu.Lock()
725	defer s.mu.Unlock()
726
727	if s.closed {
728	return nil
729	}
730	s.closed = true
731	return s.db.Close()
732	}
733
734	type scannable interface {
735	Scan(dest ...interface{}) error
736	}
737
738	func scanProbe(s scannable) (*common.ProbeResult, error) {
739	var p common.ProbeResult
740	var ts, verdStr, tsJSON string
741	err := s.Scan(
742	&p.ID, &ts, &p.Stage, &p.Protocol, &p.TargetPort, &p.TargetHost,
743	&p.LocalAddr, &p.ProxyHop, &p.TLSProfile, &p.JA3String, &p.JA4String,
744	&p.PayloadType, &p.CarrierType, &p.TargetVolume, &p.BytesSent, &p.BytesRecv,
745	&p.ThroughputBps, &p.PayloadHash, &p.WireHash, &verdStr,
746	&p.CapturedHash, &p.ErrorMessage, &p.RawEvidence, &tsJSON,
747	)
748	if err != nil {
749	if err == sql.ErrNoRows {
750	return nil, ErrProbeNotFound
751	}
752	return nil, err
753	}
754
755	p.Timestamp, _ = time.Parse(time.RFC3339Nano, ts)
756	p.Verdict = common.Verdict(verdStr)
757	if tsJSON != "" {
758	_ = json.Unmarshal([]byte(tsJSON), &p.Timestamps)
759	}
760	return &p, nil
761	}
762
763	func scanSession(s scannable, fallbackToken common.SessionToken) (*common.Session, error) {
764	var sess common.Session
765	var sessID, tokHex string
766	var tok sql.NullString
767	var sTime, lActive, verdStr string
768	err := s.Scan(
769	&sessID, &tokHex, &tok, &sess.ClientIP, &sess.Protocol, &sess.Port, &sTime, &lActive,
770	&sess.TotalBytes, &sess.ExpectedBytes, &sess.TotalChunks, &sess.ReceivedChunks,
771	&sess.PayloadHash, &sess.ExpectedHash, &sess.Status, &verdStr,
772	)
773	if err != nil {
774	if err == sql.ErrNoRows {
775	return nil, ErrSessionNotFound
776	}
777	return nil, err
778	}
779
780	if tok.Valid && tok.String != "" {
781	parsed, err := common.ParseSessionToken(tok.String)
782	if err == nil {
783	sess.Token = parsed
784	} else {
785	sess.Token = fallbackToken
786	}
787	} else if tokHex != "" {
788	parsed, err := common.ParseSessionToken(tokHex)
789	if err == nil {
790	sess.Token = parsed
791	} else {
792	sess.Token = fallbackToken
793	}
794	} else {
795	sess.Token = fallbackToken
796	}
797
798	if sessID != "" {
799	sess.SessionID = sessID
800	} else {
801	sess.SessionID = sess.Token.String()
802	}
803
804	sess.Verdict = common.Verdict(verdStr)
805	sess.StartTime, _ = time.Parse(time.RFC3339Nano, sTime)
806	sess.LastActive, _ = time.Parse(time.RFC3339Nano, lActive)
807	sess.Chunks = make(map[uint32]common.ChunkRecord)
808
809	return &sess, nil
810	}

package state

import (
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/shehzade/contraband/pkg/common"
	_ "modernc.org/sqlite" // Pure-Go SQLite driver for CGO_ENABLED=0 static compilation
)

// SQLiteLedger implements an ACID-compliant, persistent session ledger using pure-Go SQLite.
type SQLiteLedger struct {
	mu       sync.RWMutex
	db       *sql.DB
	observer LedgerObserver
	closed   bool
}

// NewSQLiteLedger opens or creates a SQLite database at the specified path.
func NewSQLiteLedger(dbPath string) (*SQLiteLedger, error) {
	db, err := sql.Open("sqlite", dbPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open sqlite database at %s: %w", dbPath, err)
	}

	db.SetMaxOpenConns(1)

	ledger := &SQLiteLedger{
		db: db,
	}

	if err := ledger.initSchema(); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("failed to initialize sqlite schema: %w", err)
	}

	return ledger, nil
}

// SetObserver attaches an observer to receive real-time chunk and session events.
func (s *SQLiteLedger) SetObserver(obs LedgerObserver) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.observer = obs
}

func (s *SQLiteLedger) initSchema() error {
	schema := `
	CREATE TABLE IF NOT EXISTS probes (
		id TEXT PRIMARY KEY,
		timestamp TEXT NOT NULL,
		stage TEXT NOT NULL,
		protocol TEXT NOT NULL,
		target_port INTEGER NOT NULL,
		target_host TEXT NOT NULL,
		local_addr TEXT,
		proxy_hop TEXT,
		tls_profile TEXT,
		ja3_string TEXT,
		ja4_string TEXT,
		payload_type TEXT,
		carrier_type TEXT,
		target_volume_bytes INTEGER,
		bytes_sent INTEGER,
		bytes_recv INTEGER,
		throughput_kbps REAL,
		payload_hash TEXT,
		wire_hash TEXT,
		verdict TEXT NOT NULL,
		captured_hash TEXT,
		error_message TEXT,
		raw_evidence TEXT,
		timestamps_json TEXT
	);

	CREATE TABLE IF NOT EXISTS sessions (
		session_id TEXT PRIMARY KEY,
		token_hex TEXT NOT NULL,
		token TEXT,
		client_ip TEXT NOT NULL,
		protocol TEXT NOT NULL,
		port INTEGER NOT NULL,
		start_time TEXT NOT NULL,
		last_active TEXT NOT NULL,
		total_bytes INTEGER NOT NULL,
		expected_bytes INTEGER NOT NULL,
		total_chunks INTEGER NOT NULL,
		received_chunks INTEGER NOT NULL,
		payload_hash TEXT,
		expected_hash TEXT,
		status TEXT NOT NULL,
		verdict TEXT NOT NULL
	);

	CREATE TABLE IF NOT EXISTS session_chunks (
		session_id TEXT NOT NULL,
		seq INTEGER NOT NULL,
		total INTEGER NOT NULL,
		data BLOB NOT NULL,
		received_at TEXT NOT NULL,
		PRIMARY KEY (session_id, seq),
		FOREIGN KEY (session_id) REFERENCES sessions(session_id) ON DELETE CASCADE
	);

	CREATE VIEW IF NOT EXISTS chunks AS
		SELECT session_id AS token, seq AS sequence, total, LENGTH(data) AS length, received_at, data AS payload
		FROM session_chunks;

	CREATE TRIGGER IF NOT EXISTS insert_chunks_trigger INSTEAD OF INSERT ON chunks
	BEGIN
		INSERT INTO session_chunks (session_id, seq, total, data, received_at)
		VALUES (NEW.token, NEW.sequence, NEW.total, NEW.payload, NEW.received_at)
		ON CONFLICT(session_id, seq) DO NOTHING;
	END;

	CREATE TABLE IF NOT EXISTS credentials (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		protocol TEXT NOT NULL,
		timestamp TEXT NOT NULL,
		client_ip TEXT NOT NULL,
		username TEXT NOT NULL,
		domain TEXT,
		hash_type TEXT NOT NULL,
		hash TEXT NOT NULL
	);

	CREATE INDEX IF NOT EXISTS idx_session_chunks_session ON session_chunks(session_id);
	CREATE INDEX IF NOT EXISTS idx_probes_stage ON probes(stage);
	CREATE INDEX IF NOT EXISTS idx_probes_verdict ON probes(verdict);
	CREATE INDEX IF NOT EXISTS idx_probes_timestamp ON probes(timestamp);
	`
	_, err := s.db.Exec(schema)
	return err
}

func (s *SQLiteLedger) RecordProbe(result common.ProbeResult) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.closed {
		return ErrLedgerClosed
	}

	if result.Timestamp.IsZero() {
		result.Timestamp = time.Now().UTC()
	}
	if result.TargetPort == 0 {
		result.TargetPort = 80
	}

	if err := result.Validate(); err != nil {
		return fmt.Errorf("invalid probe result: %w", err)
	}

	tsJSON, err := json.Marshal(result.Timestamps)
	if err != nil {
		return fmt.Errorf("failed to marshal timestamps: %w", err)
	}

	query := `
	INSERT INTO probes (
		id, timestamp, stage, protocol, target_port, target_host,
		local_addr, proxy_hop, tls_profile, ja3_string, ja4_string,
		payload_type, carrier_type, target_volume_bytes, bytes_sent, bytes_recv,
		throughput_kbps, payload_hash, wire_hash, verdict,
		captured_hash, error_message, raw_evidence, timestamps_json
	) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	ON CONFLICT(id) DO UPDATE SET
		verdict = excluded.verdict,
		bytes_sent = excluded.bytes_sent,
		bytes_recv = excluded.bytes_recv,
		throughput_kbps = excluded.throughput_kbps,
		payload_hash = excluded.payload_hash,
		wire_hash = excluded.wire_hash,
		error_message = excluded.error_message,
		raw_evidence = excluded.raw_evidence,
		timestamps_json = excluded.timestamps_json;
	`
	_, err = s.db.Exec(query,
		result.ID, result.Timestamp.Format(time.RFC3339Nano), result.Stage, result.Protocol,
		result.TargetPort, result.TargetHost, result.LocalAddr, result.ProxyHop,
		result.TLSProfile, result.JA3String, result.JA4String, result.PayloadType,
		result.CarrierType, result.TargetVolume, result.BytesSent, result.BytesRecv,
		result.ThroughputBps, result.PayloadHash, result.WireHash, string(result.Verdict),
		result.CapturedHash, result.ErrorMessage, result.RawEvidence, string(tsJSON),
	)
	if err == nil && s.observer != nil {
		if po, ok := s.observer.(ProbeObserver); ok {
			po.OnProbeRecorded(&result)
		}
	}
	return err
}

func (s *SQLiteLedger) GetProbe(id string) (*common.ProbeResult, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return nil, ErrLedgerClosed
	}

	query := `SELECT id, timestamp, stage, protocol, target_port, target_host, local_addr,
		proxy_hop, tls_profile, ja3_string, ja4_string, payload_type, carrier_type,
		target_volume_bytes, bytes_sent, bytes_recv, throughput_kbps, payload_hash,
		wire_hash, verdict, captured_hash, error_message, raw_evidence, timestamps_json
		FROM probes WHERE id = ?`

	row := s.db.QueryRow(query, id)
	return scanProbe(row)
}

func (s *SQLiteLedger) GetAllProbes() ([]common.ProbeResult, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return nil, ErrLedgerClosed
	}

	query := `SELECT id, timestamp, stage, protocol, target_port, target_host, local_addr,
		proxy_hop, tls_profile, ja3_string, ja4_string, payload_type, carrier_type,
		target_volume_bytes, bytes_sent, bytes_recv, throughput_kbps, payload_hash,
		wire_hash, verdict, captured_hash, error_message, raw_evidence, timestamps_json
		FROM probes ORDER BY timestamp ASC`

	rows, err := s.db.Query(query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var results []common.ProbeResult
	for rows.Next() {
		p, err := scanProbe(rows)
		if err != nil {
			return nil, err
		}
		results = append(results, *p)
	}
	return results, rows.Err()
}

func (s *SQLiteLedger) GetSession(token common.SessionToken) (*common.Session, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return nil, ErrLedgerClosed
	}

	tokStr := token.String()
	query := `SELECT session_id, token_hex, token, client_ip, protocol, port, start_time, last_active,
		total_bytes, expected_bytes, total_chunks, received_chunks,
		payload_hash, expected_hash, status, verdict
		FROM sessions WHERE session_id = ? OR token_hex = ? OR token = ?`

	row := s.db.QueryRow(query, tokStr, tokStr, tokStr)
	session, err := scanSession(row, token)
	if err != nil {
		return nil, err
	}

	chunkQuery := `SELECT seq, total, LENGTH(data), received_at FROM session_chunks WHERE session_id = ?`
	cRows, err := s.db.Query(chunkQuery, session.SessionID)
	if err == nil {
		defer cRows.Close()
		for cRows.Next() {
			var cr common.ChunkRecord
			var rAt string
			if err := cRows.Scan(&cr.Sequence, &cr.Total, &cr.Length, &rAt); err == nil {
				cr.ReceivedAt, _ = time.Parse(time.RFC3339Nano, rAt)
				session.Chunks[cr.Sequence] = cr
			}
		}
	}

	return session, nil
}

func (s *SQLiteLedger) UpdateSession(session *common.Session) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.closed {
		return ErrLedgerClosed
	}
	if session == nil {
		return fmt.Errorf("cannot update nil session")
	}

	sessID := session.SessionID
	tokStr := session.Token.String()
	if sessID == "" {
		sessID = tokStr
	}
	if tokStr == "" {
		tokStr = sessID
	}

	query := `INSERT INTO sessions (
		session_id, token_hex, token, client_ip, protocol, port,
		start_time, last_active, total_bytes, expected_bytes,
		total_chunks, received_chunks, payload_hash, expected_hash,
		status, verdict
	) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	ON CONFLICT(session_id) DO UPDATE SET
		status = excluded.status,
		verdict = excluded.verdict,
		last_active = excluded.last_active,
		payload_hash = CASE WHEN excluded.payload_hash != '' THEN excluded.payload_hash ELSE sessions.payload_hash END,
		expected_hash = CASE WHEN excluded.expected_hash != '' THEN excluded.expected_hash ELSE sessions.expected_hash END,
		total_bytes = CASE WHEN excluded.total_bytes > sessions.total_bytes THEN excluded.total_bytes ELSE sessions.total_bytes END,
		received_chunks = CASE WHEN excluded.received_chunks > sessions.received_chunks THEN excluded.received_chunks ELSE sessions.received_chunks END,
		total_chunks = CASE WHEN excluded.total_chunks > 0 THEN excluded.total_chunks ELSE sessions.total_chunks END;`

	_, err := s.db.Exec(query,
		sessID, tokStr, tokStr, session.ClientIP, session.Protocol, session.Port,
		session.StartTime.Format(time.RFC3339Nano),
		session.LastActive.Format(time.RFC3339Nano),
		session.TotalBytes, session.ExpectedBytes,
		session.TotalChunks, session.ReceivedChunks,
		session.PayloadHash, session.ExpectedHash,
		session.Status, string(session.Verdict),
	)
	return err
}

func (s *SQLiteLedger) GetAllSessions() ([]*common.Session, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return nil, ErrLedgerClosed
	}

	query := `SELECT session_id, token_hex, token, client_ip, protocol, port, start_time, last_active,
		total_bytes, expected_bytes, total_chunks, received_chunks,
		payload_hash, expected_hash, status, verdict
		FROM sessions ORDER BY start_time ASC`

	rows, err := s.db.Query(query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var sessions []*common.Session
	for rows.Next() {
		sess, err := scanSession(rows, common.SessionToken{})
		if err != nil {
			return nil, err
		}
		sessions = append(sessions, sess)
	}
	return sessions, rows.Err()
}

func (s *SQLiteLedger) AssembleChunk(token common.SessionToken, seq uint32, total uint32, payload []byte, expectedHash string) (*common.Session, bool, error) {
	return s.AssembleChunkWithMeta(token, seq, total, payload, expectedHash, "", "", 0, "", "")
}

func (s *SQLiteLedger) AssembleChunkWithMeta(token common.SessionToken, seq uint32, total uint32, payload []byte, expectedHash, clientIP, protocol string, port int, clientHost, networkZone string) (*common.Session, bool, error) {
	s.mu.Lock()

	if s.closed {
		s.mu.Unlock()
		return nil, false, ErrLedgerClosed
	}
	if total == 0 || seq == 0 || seq > total {
		s.mu.Unlock()
		return nil, false, fmt.Errorf("%w: seq=%d total=%d", ErrInvalidChunk, seq, total)
	}

	tokStr := token.String()
	now := time.Now().UTC()
	nowStr := now.Format(time.RFC3339Nano)

	tx, err := s.db.Begin()
	if err != nil {
		s.mu.Unlock()
		return nil, false, err
	}
	defer func() { _ = tx.Rollback() }()

	upsertSession := `
	INSERT INTO sessions (
		session_id, token_hex, token, client_ip, protocol, port, start_time, last_active,
		total_bytes, expected_bytes, total_chunks, received_chunks,
		payload_hash, expected_hash, status, verdict
	) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0, 0, ?, 0, '', ?, 'active', 'BLOCKED_L7')
	ON CONFLICT(session_id) DO UPDATE SET
		total_chunks = excluded.total_chunks,
		client_ip = CASE WHEN sessions.client_ip = '' AND excluded.client_ip != '' THEN excluded.client_ip ELSE sessions.client_ip END,
		protocol = CASE WHEN sessions.protocol = '' AND excluded.protocol != '' THEN excluded.protocol ELSE sessions.protocol END,
		port = CASE WHEN sessions.port = 0 AND excluded.port != 0 THEN excluded.port ELSE sessions.port END,
		expected_hash = CASE WHEN excluded.expected_hash != '' THEN excluded.expected_hash ELSE sessions.expected_hash END;
	`
	if _, err := tx.Exec(upsertSession, tokStr, tokStr, tokStr, clientIP, protocol, port, nowStr, nowStr, total, expectedHash); err != nil {
		s.mu.Unlock()
		return nil, false, err
	}

	insertChunk := `
	INSERT INTO session_chunks (session_id, seq, total, data, received_at)
	VALUES (?, ?, ?, ?, ?)
	ON CONFLICT(session_id, seq) DO NOTHING;
	`
	res, err := tx.Exec(insertChunk, tokStr, seq, total, payload, nowStr)
	if err != nil {
		s.mu.Unlock()
		return nil, false, err
	}

	rowsAffected, _ := res.RowsAffected()
	if rowsAffected > 0 {
		updateStats := `
		UPDATE sessions SET
			received_chunks = received_chunks + 1,
			total_bytes = total_bytes + ?,
			last_active = ?
		WHERE session_id = ?;
		`
		if _, err := tx.Exec(updateStats, len(payload), nowStr, tokStr); err != nil {
			s.mu.Unlock()
			return nil, false, err
		}
	}

	var sess common.Session
	var sessID, tokHex string
	var tokNull sql.NullString
	var sTime, lActive, verdStr string
	querySess := `SELECT session_id, token_hex, token, client_ip, protocol, port, start_time, last_active,
		total_bytes, expected_bytes, total_chunks, received_chunks,
		payload_hash, expected_hash, status, verdict
		FROM sessions WHERE session_id = ?`
	err = tx.QueryRow(querySess, tokStr).Scan(
		&sessID, &tokHex, &tokNull, &sess.ClientIP, &sess.Protocol, &sess.Port, &sTime, &lActive,
		&sess.TotalBytes, &sess.ExpectedBytes, &sess.TotalChunks, &sess.ReceivedChunks,
		&sess.PayloadHash, &sess.ExpectedHash, &sess.Status, &verdStr,
	)
	if err != nil {
		s.mu.Unlock()
		return nil, false, err
	}
	sess.Token = token
	sess.SessionID = sessID
	sess.Verdict = common.Verdict(verdStr)
	sess.StartTime, _ = time.Parse(time.RFC3339Nano, sTime)
	sess.LastActive, _ = time.Parse(time.RFC3339Nano, lActive)

	obs := s.observer

	// Idempotency: if session was already completed or corrupted, return existing state
	if sess.Status == "completed" {
		_ = tx.Rollback()
		retSess := CloneSession(&sess)
		s.mu.Unlock()
		return retSess, true, nil
	}
	if sess.Status == "corrupted" {
		_ = tx.Rollback()
		retSess := CloneSession(&sess)
		s.mu.Unlock()
		return retSess, true, ErrHashMismatch
	}

	if sess.ReceivedChunks >= int(total) {
		chunkQuery := `SELECT seq, data FROM session_chunks WHERE session_id = ? ORDER BY seq ASC`
		cRows, err := tx.Query(chunkQuery, tokStr)
		if err != nil {
			s.mu.Unlock()
			return nil, false, err
		}
		defer cRows.Close()

		chunkMap := make(map[uint32][]byte)
		for cRows.Next() {
			var sNum uint32
			var pBytes []byte
			if err := cRows.Scan(&sNum, &pBytes); err == nil {
				chunkMap[sNum] = pBytes
			}
		}

		res := ReassembleChunks(chunkMap, total, sess.ExpectedHash)
		sess.PayloadHash = res.ComputedHash
		sess.Status = res.Status
		sess.Verdict = res.Verdict

		updateFinal := `UPDATE sessions SET payload_hash = ?, status = ?, verdict = ? WHERE session_id = ?`
		if _, err := tx.Exec(updateFinal, sess.PayloadHash, sess.Status, string(sess.Verdict), tokStr); err != nil {
			s.mu.Unlock()
			return nil, false, err
		}

		if err := tx.Commit(); err != nil {
			s.mu.Unlock()
			return nil, false, err
		}

		retSess := CloneSession(&sess)
		s.mu.Unlock()

		if obs != nil {
			obs.OnChunkReceived(retSess, seq, total, int64(len(payload)))
			obs.OnSessionCompleted(retSess)
		}

		if res.Error != nil {
			return retSess, true, res.Error
		}
		return retSess, true, nil
	}

	if err := tx.Commit(); err != nil {
		s.mu.Unlock()
		return nil, false, err
	}

	retSess := CloneSession(&sess)
	s.mu.Unlock()

	if obs != nil {
		obs.OnChunkReceived(retSess, seq, total, int64(len(payload)))
	}

	return retSess, false, nil
}

func (s *SQLiteLedger) GetSessionPayload(token common.SessionToken) ([]byte, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return nil, ErrLedgerClosed
	}

	tokStr := token.String()

	var totalChunks int
	var status, verdict, expectedHash string
	err := s.db.QueryRow("SELECT total_chunks, status, verdict, expected_hash FROM sessions WHERE session_id = ? OR token_hex = ? OR token = ?", tokStr, tokStr, tokStr).Scan(&totalChunks, &status, &verdict, &expectedHash)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, ErrSessionNotFound
		}
		return nil, err
	}

	if status == "corrupted" || common.Verdict(verdict) == common.VerdictBlockedDLP {
		return nil, ErrHashMismatch
	}

	if totalChunks <= 0 {
		return nil, fmt.Errorf("%w: total chunks is %d", ErrSessionIncomplete, totalChunks)
	}

	var count int
	if err := s.db.QueryRow("SELECT COUNT(*) FROM session_chunks WHERE session_id = ? OR session_id IN (SELECT session_id FROM sessions WHERE token_hex = ? OR token = ?)", tokStr, tokStr, tokStr).Scan(&count); err != nil {
		return nil, err
	}
	if count < totalChunks {
		return nil, fmt.Errorf("%w: received %d of %d chunks", ErrSessionIncomplete, count, totalChunks)
	}

	query := `SELECT seq, data FROM session_chunks WHERE session_id = ? OR session_id IN (SELECT session_id FROM sessions WHERE token_hex = ? OR token = ?) ORDER BY seq ASC`
	rows, err := s.db.Query(query, tokStr, tokStr, tokStr)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var fullPayload []byte
	var expectedSeq uint32 = 1

	for rows.Next() {
		var seq uint32
		var payload []byte
		if err := rows.Scan(&seq, &payload); err != nil {
			return nil, err
		}
		if seq != expectedSeq {
			return nil, fmt.Errorf("%w: missing sequence %d (found %d)", ErrSessionIncomplete, expectedSeq, seq)
		}
		fullPayload = append(fullPayload, payload...)
		expectedSeq++
	}

	if int(expectedSeq-1) != totalChunks {
		return nil, fmt.Errorf("%w: missing trailing chunks: expected %d, assembled %d", ErrSessionIncomplete, totalChunks, expectedSeq-1)
	}

	if expectedHash != "" {
		hasher := sha256.New()
		hasher.Write(fullPayload)
		computed := hex.EncodeToString(hasher.Sum(nil))
		if !strings.EqualFold(computed, expectedHash) {
			return nil, ErrHashMismatch
		}
	}

	return fullPayload, nil
}

func (s *SQLiteLedger) RecordCredential(cred common.CapturedCredential) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.closed {
		return ErrLedgerClosed
	}

	query := `INSERT INTO credentials (protocol, timestamp, client_ip, username, domain, hash_type, hash)
		VALUES (?, ?, ?, ?, ?, ?, ?)`

	_, err := s.db.Exec(query,
		cred.Protocol, cred.Timestamp.Format(time.RFC3339Nano), cred.ClientIP,
		cred.Username, cred.Domain, cred.HashType, cred.Hash,
	)
	return err
}

func (s *SQLiteLedger) GetCredentials() ([]common.CapturedCredential, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return nil, ErrLedgerClosed
	}

	query := `SELECT protocol, timestamp, client_ip, username, domain, hash_type, hash FROM credentials ORDER BY id ASC`
	rows, err := s.db.Query(query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var creds []common.CapturedCredential
	for rows.Next() {
		var c common.CapturedCredential
		var ts string
		if err := rows.Scan(&c.Protocol, &ts, &c.ClientIP, &c.Username, &c.Domain, &c.HashType, &c.Hash); err != nil {
			return nil, err
		}
		c.Timestamp, _ = time.Parse(time.RFC3339Nano, ts)
		creds = append(creds, c)
	}
	return creds, rows.Err()
}

func (s *SQLiteLedger) ExportJSON() ([]byte, error) {
	probes, err := s.GetAllProbes()
	if err != nil {
		return nil, err
	}
	sessions, err := s.GetAllSessions()
	if err != nil {
		return nil, err
	}
	creds, err := s.GetCredentials()
	if err != nil {
		return nil, err
	}

	dump := struct {
		TotalProbes      int                         `json:"total_probes"`
		TotalSessions    int                         `json:"total_sessions"`
		TotalCredentials int                         `json:"total_credentials"`
		Probes           []common.ProbeResult        `json:"probes"`
		Sessions         []*common.Session           `json:"sessions"`
		Credentials      []common.CapturedCredential `json:"credentials"`
	}{
		TotalProbes:      len(probes),
		TotalSessions:    len(sessions),
		TotalCredentials: len(creds),
		Probes:           probes,
		Sessions:         sessions,
		Credentials:      creds,
	}

	return json.MarshalIndent(dump, "", "  ")
}

// Reset clears all records from session_chunks, sessions, probes, and credentials tables.
func (s *SQLiteLedger) Reset() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.closed {
		return ErrLedgerClosed
	}

	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	if _, err := tx.Exec(`DELETE FROM session_chunks`); err != nil {
		return err
	}
	if _, err := tx.Exec(`DELETE FROM sessions`); err != nil {
		return err
	}
	if _, err := tx.Exec(`DELETE FROM probes`); err != nil {
		return err
	}
	if _, err := tx.Exec(`DELETE FROM credentials`); err != nil {
		return err
	}

	return tx.Commit()
}

func (s *SQLiteLedger) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.closed {
		return nil
	}
	s.closed = true
	return s.db.Close()
}

type scannable interface {
	Scan(dest ...interface{}) error
}

func scanProbe(s scannable) (*common.ProbeResult, error) {
	var p common.ProbeResult
	var ts, verdStr, tsJSON string
	err := s.Scan(
		&p.ID, &ts, &p.Stage, &p.Protocol, &p.TargetPort, &p.TargetHost,
		&p.LocalAddr, &p.ProxyHop, &p.TLSProfile, &p.JA3String, &p.JA4String,
		&p.PayloadType, &p.CarrierType, &p.TargetVolume, &p.BytesSent, &p.BytesRecv,
		&p.ThroughputBps, &p.PayloadHash, &p.WireHash, &verdStr,
		&p.CapturedHash, &p.ErrorMessage, &p.RawEvidence, &tsJSON,
	)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, ErrProbeNotFound
		}
		return nil, err
	}

	p.Timestamp, _ = time.Parse(time.RFC3339Nano, ts)
	p.Verdict = common.Verdict(verdStr)
	if tsJSON != "" {
		_ = json.Unmarshal([]byte(tsJSON), &p.Timestamps)
	}
	return &p, nil
}

func scanSession(s scannable, fallbackToken common.SessionToken) (*common.Session, error) {
	var sess common.Session
	var sessID, tokHex string
	var tok sql.NullString
	var sTime, lActive, verdStr string
	err := s.Scan(
		&sessID, &tokHex, &tok, &sess.ClientIP, &sess.Protocol, &sess.Port, &sTime, &lActive,
		&sess.TotalBytes, &sess.ExpectedBytes, &sess.TotalChunks, &sess.ReceivedChunks,
		&sess.PayloadHash, &sess.ExpectedHash, &sess.Status, &verdStr,
	)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, ErrSessionNotFound
		}
		return nil, err
	}

	if tok.Valid && tok.String != "" {
		parsed, err := common.ParseSessionToken(tok.String)
		if err == nil {
			sess.Token = parsed
		} else {
			sess.Token = fallbackToken
		}
	} else if tokHex != "" {
		parsed, err := common.ParseSessionToken(tokHex)
		if err == nil {
			sess.Token = parsed
		} else {
			sess.Token = fallbackToken
		}
	} else {
		sess.Token = fallbackToken
	}

	if sessID != "" {
		sess.SessionID = sessID
	} else {
		sess.SessionID = sess.Token.String()
	}

	sess.Verdict = common.Verdict(verdStr)
	sess.StartTime, _ = time.Parse(time.RFC3339Nano, sTime)
	sess.LastActive, _ = time.Parse(time.RFC3339Nano, lActive)
	sess.Chunks = make(map[uint32]common.ChunkRecord)

	return &sess, nil
}