pkg/server/state/memory_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

693 lines · 18.4 KB

Raw

1	package state
2
3	import (
4	"crypto/sha256"
5	"encoding/hex"
6	"encoding/json"
7	"fmt"
8	"io"
9	"os"
10	"strings"
11	"sync"
12	"time"
13
14	"github.com/shehzade/contraband/pkg/common"
15	)
16
17	// MemoryLedger implements an in-memory, thread-safe session ledger and chunk assembler.
18	type MemoryLedger struct {
19	mu sync.RWMutex
20	probes map[string]common.ProbeResult
21	sessions map[common.SessionToken]*common.Session
22	chunks map[common.SessionToken]map[uint32][]byte
23	chunkSizes map[common.SessionToken]int64
24	diskFiles map[common.SessionToken]*os.File
25	diskPaths map[common.SessionToken]string
26	completedTokens []common.SessionToken
27	credentials []common.CapturedCredential
28	observer LedgerObserver
29	closed bool
30	}
31
32	// NewMemoryLedger constructs an empty, initialized in-memory ledger.
33	func NewMemoryLedger() *MemoryLedger {
34	return &MemoryLedger{
35	probes: make(map[string]common.ProbeResult),
36	sessions: make(map[common.SessionToken]*common.Session),
37	chunks: make(map[common.SessionToken]map[uint32][]byte),
38	chunkSizes: make(map[common.SessionToken]int64),
39	diskFiles: make(map[common.SessionToken]*os.File),
40	diskPaths: make(map[common.SessionToken]string),
41	completedTokens: make([]common.SessionToken, 0, 64),
42	credentials: make([]common.CapturedCredential, 0),
43	}
44	}
45
46	// SetObserver attaches an observer to receive real-time chunk and session events.
47	func (m *MemoryLedger) SetObserver(obs LedgerObserver) {
48	m.mu.Lock()
49	defer m.mu.Unlock()
50	m.observer = obs
51	}
52
53	// RecordProbe stores or updates an individual probe result with strict validation.
54	func (m *MemoryLedger) RecordProbe(result common.ProbeResult) error {
55	m.mu.Lock()
56	defer m.mu.Unlock()
57
58	if m.closed {
59	return ErrLedgerClosed
60	}
61
62	if result.Timestamp.IsZero() {
63	result.Timestamp = time.Now().UTC()
64	}
65	if result.TargetPort == 0 {
66	result.TargetPort = 80
67	}
68
69	if err := result.Validate(); err != nil {
70	return fmt.Errorf("invalid probe result: %w", err)
71	}
72
73	m.probes[result.ID] = result
74	if m.observer != nil {
75	if po, ok := m.observer.(ProbeObserver); ok {
76	po.OnProbeRecorded(&result)
77	}
78	}
79	return nil
80	}
81
82	// GetProbe retrieves a probe by its unique UUID v4 string.
83	func (m MemoryLedger) GetProbe(id string) (common.ProbeResult, error) {
84	m.mu.RLock()
85	defer m.mu.RUnlock()
86
87	if m.closed {
88	return nil, ErrLedgerClosed
89	}
90
91	p, exists := m.probes[id]
92	if !exists {
93	return nil, ErrProbeNotFound
94	}
95
96	copied := p
97	return &copied, nil
98	}
99
100	// GetAllProbes returns a snapshot slice of all recorded probe results.
101	func (m *MemoryLedger) GetAllProbes() ([]common.ProbeResult, error) {
102	m.mu.RLock()
103	defer m.mu.RUnlock()
104
105	if m.closed {
106	return nil, ErrLedgerClosed
107	}
108
109	results := make([]common.ProbeResult, 0, len(m.probes))
110	for _, p := range m.probes {
111	results = append(results, p)
112	}
113	return results, nil
114	}
115
116	// GetSession retrieves a session by its 16-byte session token.
117	func (m MemoryLedger) GetSession(token common.SessionToken) (common.Session, error) {
118	m.mu.RLock()
119	defer m.mu.RUnlock()
120
121	if m.closed {
122	return nil, ErrLedgerClosed
123	}
124
125	session, exists := m.sessions[token]
126	if !exists {
127	return nil, ErrSessionNotFound
128	}
129
130	return CloneSession(session), nil
131	}
132
133	// UpdateSession updates the mutable fields of an active session, creating it if it does not exist.
134	func (m MemoryLedger) UpdateSession(session common.Session) error {
135	m.mu.Lock()
136	defer m.mu.Unlock()
137
138	if m.closed {
139	return ErrLedgerClosed
140	}
141	if session == nil {
142	return fmt.Errorf("cannot update nil session")
143	}
144
145	existing, exists := m.sessions[session.Token]
146	if !exists {
147	m.sessions[session.Token] = CloneSession(session)
148	if _, ok := m.chunks[session.Token]; !ok {
149	m.chunks[session.Token] = make(map[uint32][]byte)
150	}
151	return nil
152	}
153
154	if session.ClientIP != "" {
155	existing.ClientIP = session.ClientIP
156	}
157	if session.Protocol != "" {
158	existing.Protocol = session.Protocol
159	}
160	if session.Port != 0 {
161	existing.Port = session.Port
162	}
163	existing.Status = session.Status
164	existing.Verdict = session.Verdict
165	if !session.LastActive.IsZero() {
166	existing.LastActive = session.LastActive
167	}
168	if session.TotalBytes != 0 {
169	existing.TotalBytes = session.TotalBytes
170	}
171	if session.ExpectedBytes != 0 {
172	existing.ExpectedBytes = session.ExpectedBytes
173	}
174	if session.TotalChunks != 0 {
175	existing.TotalChunks = session.TotalChunks
176	}
177	if session.ReceivedChunks != 0 {
178	existing.ReceivedChunks = session.ReceivedChunks
179	}
180	existing.PayloadHash = session.PayloadHash
181	if session.ExpectedHash != "" {
182	existing.ExpectedHash = session.ExpectedHash
183	}
184
185	return nil
186	}
187
188	// GetAllSessions returns a snapshot of all active and completed sessions.
189	func (m MemoryLedger) GetAllSessions() ([]common.Session, error) {
190	m.mu.RLock()
191	defer m.mu.RUnlock()
192
193	if m.closed {
194	return nil, ErrLedgerClosed
195	}
196
197	sessions := make([]*common.Session, 0, len(m.sessions))
198	for _, s := range m.sessions {
199	sessions = append(sessions, CloneSession(s))
200	}
201	return sessions, nil
202	}
203
204	// AssembleChunk buffers an incoming chunk, updates session metrics, triggers reassembly upon completion,
205	// and notifies any attached observer.
206	func (m MemoryLedger) AssembleChunk(token common.SessionToken, seq uint32, total uint32, payload []byte, expectedHash string) (common.Session, bool, error) {
207	return m.AssembleChunkWithMeta(token, seq, total, payload, expectedHash, "", "", 0, "", "")
208	}
209
210	// AssembleChunkWithMeta buffers an incoming chunk with explicit client IP, protocol, port, host, and network zone metadata.
211	func (m MemoryLedger) AssembleChunkWithMeta(token common.SessionToken, seq uint32, total uint32, payload []byte, expectedHash, clientIP, protocol string, port int, clientHost, networkZone string) (common.Session, bool, error) {
212	m.mu.Lock()
213
214	if m.closed {
215	m.mu.Unlock()
216	return nil, false, ErrLedgerClosed
217	}
218	if total == 0 \|\| seq == 0 \|\| seq > total {
219	m.mu.Unlock()
220	return nil, false, fmt.Errorf("%w: seq=%d total=%d", ErrInvalidChunk, seq, total)
221	}
222
223	session, exists := m.sessions[token]
224	if !exists {
225	session = common.NewSession(token, clientIP, protocol, port)
226	session.ClientHost = clientHost
227	session.NetworkZone = networkZone
228	session.TotalChunks = int(total)
229	session.ExpectedHash = expectedHash
230	m.sessions[token] = session
231	m.chunks[token] = make(map[uint32][]byte)
232	} else {
233	// Idempotency: if session was already completed or corrupted, do not allow stray/duplicate
234	// chunks to overwrite state or re-trigger reassembly.
235	if session.Status == "completed" {
236	retSession := CloneSession(session)
237	m.mu.Unlock()
238	return retSession, true, nil
239	}
240	if session.Status == "corrupted" {
241	retSession := CloneSession(session)
242	m.mu.Unlock()
243	return retSession, true, ErrHashMismatch
244	}
245
246	if session.ClientIP == "" && clientIP != "" {
247	session.ClientIP = clientIP
248	}
249	if session.Protocol == "" && protocol != "" {
250	session.Protocol = protocol
251	}
252	if session.Port == 0 && port > 0 {
253	session.Port = port
254	}
255	if session.ClientHost == "" && clientHost != "" {
256	session.ClientHost = clientHost
257	}
258	if session.NetworkZone == "" && networkZone != "" {
259	session.NetworkZone = networkZone
260	}
261	}
262
263	// Dynamic disk spill: if cumulative session bytes exceed 2MB or chunk length >= 1MB,
264	// spill to a temporary file so resident heap stays strictly < 32MB invariant.
265	isSpilled := false
266	var diskF *os.File
267	if m.diskFiles != nil {
268	diskF = m.diskFiles[token]
269	if diskF != nil {
270	isSpilled = true
271	}
272	}
273
274	if !isSpilled && (session.TotalBytes+int64(len(payload)) > 210241024 \|\| len(payload) >= 1024*1024) {
275	tmpF, err := os.CreateTemp("", fmt.Sprintf("ctrb_%s_*.dat", token.String()))
276	if err == nil {
277	m.diskFiles[token] = tmpF
278	diskF = tmpF
279	isSpilled = true
280
281	// Determine canonical chunk size
282	cSize := m.chunkSizes[token]
283	if cSize <= 0 {
284	if seq < total {
285	cSize = int64(len(payload))
286	} else if total == 1 {
287	cSize = int64(len(payload))
288	} else if existingChunks, ok := m.chunks[token]; ok {
289	for sIdx, cData := range existingChunks {
290	if sIdx < total && len(cData) > 0 {
291	cSize = int64(len(cData))
292	break
293	}
294	}
295	}
296	if cSize > 0 {
297	m.chunkSizes[token] = cSize
298	}
299	}
300
301	// Flush existing in-memory chunks to temp file if cSize is known
302	if cSize > 0 {
303	if existingChunks, ok := m.chunks[token]; ok {
304	for sIdx, cData := range existingChunks {
305	if len(cData) > 0 {
306	offset := int64(sIdx-1) * cSize
307	_, _ = tmpF.WriteAt(cData, offset)
308	existingChunks[sIdx] = nil // free memory, retain sequence key for deduplication
309	}
310	}
311	}
312	}
313	}
314	}
315
316	if isSpilled && diskF != nil {
317	tokenChunks := m.chunks[token]
318	if tokenChunks == nil {
319	tokenChunks = make(map[uint32][]byte)
320	m.chunks[token] = tokenChunks
321	}
322
323	// Sequence Deduplication: if this sequence has already been received, treat as idempotent duplicate.
324	// Avoids duplicate ReceivedChunks increment and premature reassembly on incomplete disk file.
325	if _, duplicate := tokenChunks[seq]; duplicate {
326	retSession := CloneSession(session)
327	m.mu.Unlock()
328	return retSession, false, nil
329	}
330
331	cSize := m.chunkSizes[token]
332	if cSize <= 0 {
333	if seq < total {
334	cSize = int64(len(payload))
335	} else if total == 1 {
336	cSize = int64(len(payload))
337	} else if existingChunks, ok := m.chunks[token]; ok {
338	for sIdx, cData := range existingChunks {
339	if sIdx < total && len(cData) > 0 {
340	cSize = int64(len(cData))
341	break
342	}
343	}
344	}
345
346	if cSize > 0 {
347	m.chunkSizes[token] = cSize
348	// Canonical chunk size just discovered; flush any previously buffered chunks (e.g. remainder chunk seq == total)
349	for sIdx, cData := range tokenChunks {
350	if len(cData) > 0 {
351	remOffset := int64(sIdx-1) * cSize
352	_, _ = diskF.WriteAt(cData, remOffset)
353	tokenChunks[sIdx] = nil
354	}
355	}
356	}
357	}
358
359	if cSize > 0 {
360	offset := int64(seq-1) * cSize
361	_, _ = diskF.WriteAt(payload, offset)
362	tokenChunks[seq] = nil
363	} else {
364	// Canonical chunk size not yet known (remainder chunk arrived first for total > 1).
365	// Buffer chunk in memory until canonical size is determined.
366	chunkCopy := make([]byte, len(payload))
367	copy(chunkCopy, payload)
368	tokenChunks[seq] = chunkCopy
369	}
370
371	session.ReceivedChunks = len(tokenChunks)
372	session.TotalBytes += int64(len(payload))
373	session.LastActive = time.Now().UTC()
374	if len(session.Chunks) < 32 {
375	session.Chunks[seq] = common.ChunkRecord{
376	Sequence: seq,
377	Total: total,
378	Length: uint32(len(payload)),
379	ReceivedAt: time.Now().UTC(),
380	}
381	}
382	} else {
383	tokenChunks := m.chunks[token]
384	if _, duplicate := tokenChunks[seq]; !duplicate {
385	chunkCopy := make([]byte, len(payload))
386	copy(chunkCopy, payload)
387	tokenChunks[seq] = chunkCopy
388
389	session.ReceivedChunks = len(tokenChunks)
390	session.TotalBytes += int64(len(payload))
391	session.LastActive = time.Now().UTC()
392	if len(session.Chunks) < 32 {
393	session.Chunks[seq] = common.ChunkRecord{
394	Sequence: seq,
395	Total: total,
396	Length: uint32(len(payload)),
397	ReceivedAt: time.Now().UTC(),
398	}
399	}
400	}
401	}
402
403	obs := m.observer
404
405	// Check if all chunks have arrived
406	if session.ReceivedChunks >= int(total) {
407	var reasmErr error
408	if isSpilled && diskF != nil {
409	// Ensure any buffered chunk (e.g. remainder chunk) is flushed to disk
410	if cSize := m.chunkSizes[token]; cSize > 0 {
411	if tokenChunks, ok := m.chunks[token]; ok {
412	for sIdx, cData := range tokenChunks {
413	if len(cData) > 0 {
414	offset := int64(sIdx-1) * cSize
415	_, _ = diskF.WriteAt(cData, offset)
416	tokenChunks[sIdx] = nil
417	}
418	}
419	}
420	}
421	_ = diskF.Sync()
422	_, _ = diskF.Seek(0, 0)
423	hasher := sha256.New()
424	buf := make([]byte, 64*1024)
425	_, _ = io.CopyBuffer(hasher, diskF, buf)
426	computedHash := hex.EncodeToString(hasher.Sum(nil))
427
428	if session.ExpectedHash != "" && !strings.EqualFold(computedHash, session.ExpectedHash) {
429	session.PayloadHash = computedHash
430	session.Status = "corrupted"
431	session.Verdict = common.VerdictBlockedDLP
432	reasmErr = ErrHashMismatch
433	} else {
434	session.PayloadHash = computedHash
435	session.Status = "completed"
436	session.Verdict = common.VerdictSuccess
437	}
438	tmpPath := diskF.Name()
439	_ = diskF.Close()
440	delete(m.diskFiles, token)
441	delete(m.chunks, token)
442	delete(m.chunkSizes, token)
443	if m.diskPaths != nil {
444	m.diskPaths[token] = tmpPath
445	}
446	m.completedTokens = append(m.completedTokens, token)
447	if len(m.completedTokens) > 2000 {
448	oldTok := m.completedTokens[0]
449	m.completedTokens = m.completedTokens[1:]
450	delete(m.chunks, oldTok)
451	if m.diskPaths != nil {
452	if oldPath, ok := m.diskPaths[oldTok]; ok {
453	_ = os.Remove(oldPath)
454	delete(m.diskPaths, oldTok)
455	}
456	}
457	}
458	} else {
459	tokenChunks := m.chunks[token]
460	res := ReassembleChunks(tokenChunks, total, session.ExpectedHash)
461	session.PayloadHash = res.ComputedHash
462	session.Status = res.Status
463	session.Verdict = res.Verdict
464
465	m.completedTokens = append(m.completedTokens, token)
466	if len(m.completedTokens) > 2000 {
467	oldTok := m.completedTokens[0]
468	m.completedTokens = m.completedTokens[1:]
469	delete(m.chunks, oldTok)
470	if m.diskPaths != nil {
471	if oldPath, ok := m.diskPaths[oldTok]; ok {
472	_ = os.Remove(oldPath)
473	delete(m.diskPaths, oldTok)
474	}
475	}
476	}
477
478	if res.Error != nil {
479	reasmErr = res.Error
480	}
481	}
482
483	if reasmErr != nil {
484	retSession := CloneSession(session)
485	m.mu.Unlock()
486	if obs != nil {
487	obs.OnChunkReceived(retSession, seq, total, int64(len(payload)))
488	obs.OnSessionCompleted(retSession)
489	}
490	return retSession, true, reasmErr
491	}
492
493	retSession := CloneSession(session)
494	m.mu.Unlock()
495
496	if obs != nil {
497	obs.OnChunkReceived(retSession, seq, total, int64(len(payload)))
498	obs.OnSessionCompleted(retSession)
499	}
500
501	return retSession, true, nil
502	}
503
504	session.Status = "active"
505	session.Verdict = common.VerdictBlockedL7
506	retSession := CloneSession(session)
507	m.mu.Unlock()
508
509	if obs != nil {
510	obs.OnChunkReceived(retSession, seq, total, int64(len(payload)))
511	}
512
513	return retSession, false, nil
514	}
515
516	// GetSessionPayload reconstructs the full contiguous payload for a session.
517	func (m *MemoryLedger) GetSessionPayload(token common.SessionToken) ([]byte, error) {
518	m.mu.RLock()
519	defer m.mu.RUnlock()
520
521	if m.closed {
522	return nil, ErrLedgerClosed
523	}
524
525	session, exists := m.sessions[token]
526	if !exists {
527	return nil, ErrSessionNotFound
528	}
529
530	if session.Status == "corrupted" \|\| session.Verdict == common.VerdictBlockedDLP {
531	return nil, ErrHashMismatch
532	}
533
534	if m.diskPaths != nil {
535	if p, ok := m.diskPaths[token]; ok && p != "" {
536	return os.ReadFile(p)
537	}
538	}
539
540	if m.diskFiles != nil {
541	if df, ok := m.diskFiles[token]; ok && df != nil {
542	_ = df.Sync()
543	_, _ = df.Seek(0, 0)
544	return io.ReadAll(df)
545	}
546	}
547
548	chunks, exists := m.chunks[token]
549	if !exists \|\| len(chunks) == 0 {
550	return nil, ErrSessionIncomplete
551	}
552
553	total := uint32(session.TotalChunks)
554	if total == 0 {
555	return nil, fmt.Errorf("%w: total chunks is 0", ErrSessionIncomplete)
556	}
557	if len(chunks) < int(total) {
558	return nil, fmt.Errorf("%w: received %d of %d chunks", ErrSessionIncomplete, len(chunks), total)
559	}
560
561	var fullPayload []byte
562	for s := uint32(1); s <= total; s++ {
563	chunkData, found := chunks[s]
564	if !found {
565	return nil, fmt.Errorf("%w: missing chunk sequence %d", ErrSessionIncomplete, s)
566	}
567	fullPayload = append(fullPayload, chunkData...)
568	}
569
570	if session.ExpectedHash != "" {
571	hasher := sha256.New()
572	hasher.Write(fullPayload)
573	computed := hex.EncodeToString(hasher.Sum(nil))
574	if !strings.EqualFold(computed, session.ExpectedHash) {
575	return nil, ErrHashMismatch
576	}
577	}
578
579	return fullPayload, nil
580	}
581
582	// RecordCredential stores harvested credentials in memory.
583	func (m *MemoryLedger) RecordCredential(cred common.CapturedCredential) error {
584	m.mu.Lock()
585	defer m.mu.Unlock()
586
587	if m.closed {
588	return ErrLedgerClosed
589	}
590
591	m.credentials = append(m.credentials, cred)
592	return nil
593	}
594
595	// GetCredentials returns a snapshot of all captured credentials.
596	func (m *MemoryLedger) GetCredentials() ([]common.CapturedCredential, error) {
597	m.mu.RLock()
598	defer m.mu.RUnlock()
599
600	if m.closed {
601	return nil, ErrLedgerClosed
602	}
603
604	copied := make([]common.CapturedCredential, len(m.credentials))
605	copy(copied, m.credentials)
606	return copied, nil
607	}
608
609	// ExportJSON serializes the entire ledger state into JSON.
610	func (m *MemoryLedger) ExportJSON() ([]byte, error) {
611	m.mu.RLock()
612	defer m.mu.RUnlock()
613
614	if m.closed {
615	return nil, ErrLedgerClosed
616	}
617
618	probes := make([]common.ProbeResult, 0, len(m.probes))
619	for _, p := range m.probes {
620	probes = append(probes, p)
621	}
622
623	sessions := make([]*common.Session, 0, len(m.sessions))
624	for _, s := range m.sessions {
625	sessions = append(sessions, CloneSession(s))
626	}
627
628	dump := struct {
629	TotalProbes int `json:"total_probes"`
630	TotalSessions int `json:"total_sessions"`
631	TotalCredentials int `json:"total_credentials"`
632	Probes []common.ProbeResult `json:"probes"`
633	Sessions []*common.Session `json:"sessions"`
634	Credentials []common.CapturedCredential `json:"credentials"`
635	}{
636	TotalProbes: len(probes),
637	TotalSessions: len(sessions),
638	TotalCredentials: len(m.credentials),
639	Probes: probes,
640	Sessions: sessions,
641	Credentials: m.credentials,
642	}
643
644	return json.MarshalIndent(dump, "", " ")
645	}
646
647	// Reset clears all recorded probes, sessions, chunks, credentials, and disk buffer files.
648	func (m *MemoryLedger) Reset() error {
649	m.mu.Lock()
650	defer m.mu.Unlock()
651
652	if m.closed {
653	return ErrLedgerClosed
654	}
655
656	for tok, f := range m.diskFiles {
657	_ = f.Close()
658	_ = os.Remove(f.Name())
659	delete(m.diskFiles, tok)
660	}
661	for tok, p := range m.diskPaths {
662	_ = os.Remove(p)
663	delete(m.diskPaths, tok)
664	}
665
666	m.probes = make(map[string]common.ProbeResult)
667	m.sessions = make(map[common.SessionToken]*common.Session)
668	m.chunks = make(map[common.SessionToken]map[uint32][]byte)
669	m.chunkSizes = make(map[common.SessionToken]int64)
670	m.completedTokens = make([]common.SessionToken, 0, 64)
671	m.credentials = make([]common.CapturedCredential, 0)
672	return nil
673	}
674
675	// Close marks the ledger as closed and purges stored chunks.
676	func (m *MemoryLedger) Close() error {
677	m.mu.Lock()
678	defer m.mu.Unlock()
679
680	m.closed = true
681	m.chunks = make(map[common.SessionToken]map[uint32][]byte)
682	m.chunkSizes = make(map[common.SessionToken]int64)
683	for tok, f := range m.diskFiles {
684	_ = f.Close()
685	_ = os.Remove(f.Name())
686	delete(m.diskFiles, tok)
687	}
688	for tok, p := range m.diskPaths {
689	_ = os.Remove(p)
690	delete(m.diskPaths, tok)
691	}
692	return nil
693	}

package state

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/shehzade/contraband/pkg/common"
)

// MemoryLedger implements an in-memory, thread-safe session ledger and chunk assembler.
type MemoryLedger struct {
	mu              sync.RWMutex
	probes          map[string]common.ProbeResult
	sessions        map[common.SessionToken]*common.Session
	chunks          map[common.SessionToken]map[uint32][]byte
	chunkSizes      map[common.SessionToken]int64
	diskFiles       map[common.SessionToken]*os.File
	diskPaths       map[common.SessionToken]string
	completedTokens []common.SessionToken
	credentials     []common.CapturedCredential
	observer        LedgerObserver
	closed          bool
}

// NewMemoryLedger constructs an empty, initialized in-memory ledger.
func NewMemoryLedger() *MemoryLedger {
	return &MemoryLedger{
		probes:          make(map[string]common.ProbeResult),
		sessions:        make(map[common.SessionToken]*common.Session),
		chunks:          make(map[common.SessionToken]map[uint32][]byte),
		chunkSizes:      make(map[common.SessionToken]int64),
		diskFiles:       make(map[common.SessionToken]*os.File),
		diskPaths:       make(map[common.SessionToken]string),
		completedTokens: make([]common.SessionToken, 0, 64),
		credentials:     make([]common.CapturedCredential, 0),
	}
}

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

// RecordProbe stores or updates an individual probe result with strict validation.
func (m *MemoryLedger) RecordProbe(result common.ProbeResult) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.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)
	}

	m.probes[result.ID] = result
	if m.observer != nil {
		if po, ok := m.observer.(ProbeObserver); ok {
			po.OnProbeRecorded(&result)
		}
	}
	return nil
}

// GetProbe retrieves a probe by its unique UUID v4 string.
func (m *MemoryLedger) GetProbe(id string) (*common.ProbeResult, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	if m.closed {
		return nil, ErrLedgerClosed
	}

	p, exists := m.probes[id]
	if !exists {
		return nil, ErrProbeNotFound
	}

	copied := p
	return &copied, nil
}

// GetAllProbes returns a snapshot slice of all recorded probe results.
func (m *MemoryLedger) GetAllProbes() ([]common.ProbeResult, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	if m.closed {
		return nil, ErrLedgerClosed
	}

	results := make([]common.ProbeResult, 0, len(m.probes))
	for _, p := range m.probes {
		results = append(results, p)
	}
	return results, nil
}

// GetSession retrieves a session by its 16-byte session token.
func (m *MemoryLedger) GetSession(token common.SessionToken) (*common.Session, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	if m.closed {
		return nil, ErrLedgerClosed
	}

	session, exists := m.sessions[token]
	if !exists {
		return nil, ErrSessionNotFound
	}

	return CloneSession(session), nil
}

// UpdateSession updates the mutable fields of an active session, creating it if it does not exist.
func (m *MemoryLedger) UpdateSession(session *common.Session) error {
	m.mu.Lock()
	defer m.mu.Unlock()

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

	existing, exists := m.sessions[session.Token]
	if !exists {
		m.sessions[session.Token] = CloneSession(session)
		if _, ok := m.chunks[session.Token]; !ok {
			m.chunks[session.Token] = make(map[uint32][]byte)
		}
		return nil
	}

	if session.ClientIP != "" {
		existing.ClientIP = session.ClientIP
	}
	if session.Protocol != "" {
		existing.Protocol = session.Protocol
	}
	if session.Port != 0 {
		existing.Port = session.Port
	}
	existing.Status = session.Status
	existing.Verdict = session.Verdict
	if !session.LastActive.IsZero() {
		existing.LastActive = session.LastActive
	}
	if session.TotalBytes != 0 {
		existing.TotalBytes = session.TotalBytes
	}
	if session.ExpectedBytes != 0 {
		existing.ExpectedBytes = session.ExpectedBytes
	}
	if session.TotalChunks != 0 {
		existing.TotalChunks = session.TotalChunks
	}
	if session.ReceivedChunks != 0 {
		existing.ReceivedChunks = session.ReceivedChunks
	}
	existing.PayloadHash = session.PayloadHash
	if session.ExpectedHash != "" {
		existing.ExpectedHash = session.ExpectedHash
	}

	return nil
}

// GetAllSessions returns a snapshot of all active and completed sessions.
func (m *MemoryLedger) GetAllSessions() ([]*common.Session, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	if m.closed {
		return nil, ErrLedgerClosed
	}

	sessions := make([]*common.Session, 0, len(m.sessions))
	for _, s := range m.sessions {
		sessions = append(sessions, CloneSession(s))
	}
	return sessions, nil
}

// AssembleChunk buffers an incoming chunk, updates session metrics, triggers reassembly upon completion,
// and notifies any attached observer.
func (m *MemoryLedger) AssembleChunk(token common.SessionToken, seq uint32, total uint32, payload []byte, expectedHash string) (*common.Session, bool, error) {
	return m.AssembleChunkWithMeta(token, seq, total, payload, expectedHash, "", "", 0, "", "")
}

// AssembleChunkWithMeta buffers an incoming chunk with explicit client IP, protocol, port, host, and network zone metadata.
func (m *MemoryLedger) AssembleChunkWithMeta(token common.SessionToken, seq uint32, total uint32, payload []byte, expectedHash, clientIP, protocol string, port int, clientHost, networkZone string) (*common.Session, bool, error) {
	m.mu.Lock()

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

	session, exists := m.sessions[token]
	if !exists {
		session = common.NewSession(token, clientIP, protocol, port)
		session.ClientHost = clientHost
		session.NetworkZone = networkZone
		session.TotalChunks = int(total)
		session.ExpectedHash = expectedHash
		m.sessions[token] = session
		m.chunks[token] = make(map[uint32][]byte)
	} else {
		// Idempotency: if session was already completed or corrupted, do not allow stray/duplicate
		// chunks to overwrite state or re-trigger reassembly.
		if session.Status == "completed" {
			retSession := CloneSession(session)
			m.mu.Unlock()
			return retSession, true, nil
		}
		if session.Status == "corrupted" {
			retSession := CloneSession(session)
			m.mu.Unlock()
			return retSession, true, ErrHashMismatch
		}

		if session.ClientIP == "" && clientIP != "" {
			session.ClientIP = clientIP
		}
		if session.Protocol == "" && protocol != "" {
			session.Protocol = protocol
		}
		if session.Port == 0 && port > 0 {
			session.Port = port
		}
		if session.ClientHost == "" && clientHost != "" {
			session.ClientHost = clientHost
		}
		if session.NetworkZone == "" && networkZone != "" {
			session.NetworkZone = networkZone
		}
	}

	// Dynamic disk spill: if cumulative session bytes exceed 2MB or chunk length >= 1MB,
	// spill to a temporary file so resident heap stays strictly < 32MB invariant.
	isSpilled := false
	var diskF *os.File
	if m.diskFiles != nil {
		diskF = m.diskFiles[token]
		if diskF != nil {
			isSpilled = true
		}
	}

	if !isSpilled && (session.TotalBytes+int64(len(payload)) > 2*1024*1024 || len(payload) >= 1024*1024) {
		tmpF, err := os.CreateTemp("", fmt.Sprintf("ctrb_%s_*.dat", token.String()))
		if err == nil {
			m.diskFiles[token] = tmpF
			diskF = tmpF
			isSpilled = true

			// Determine canonical chunk size
			cSize := m.chunkSizes[token]
			if cSize <= 0 {
				if seq < total {
					cSize = int64(len(payload))
				} else if total == 1 {
					cSize = int64(len(payload))
				} else if existingChunks, ok := m.chunks[token]; ok {
					for sIdx, cData := range existingChunks {
						if sIdx < total && len(cData) > 0 {
							cSize = int64(len(cData))
							break
						}
					}
				}
				if cSize > 0 {
					m.chunkSizes[token] = cSize
				}
			}

			// Flush existing in-memory chunks to temp file if cSize is known
			if cSize > 0 {
				if existingChunks, ok := m.chunks[token]; ok {
					for sIdx, cData := range existingChunks {
						if len(cData) > 0 {
							offset := int64(sIdx-1) * cSize
							_, _ = tmpF.WriteAt(cData, offset)
							existingChunks[sIdx] = nil // free memory, retain sequence key for deduplication
						}
					}
				}
			}
		}
	}

	if isSpilled && diskF != nil {
		tokenChunks := m.chunks[token]
		if tokenChunks == nil {
			tokenChunks = make(map[uint32][]byte)
			m.chunks[token] = tokenChunks
		}

		// Sequence Deduplication: if this sequence has already been received, treat as idempotent duplicate.
		// Avoids duplicate ReceivedChunks increment and premature reassembly on incomplete disk file.
		if _, duplicate := tokenChunks[seq]; duplicate {
			retSession := CloneSession(session)
			m.mu.Unlock()
			return retSession, false, nil
		}

		cSize := m.chunkSizes[token]
		if cSize <= 0 {
			if seq < total {
				cSize = int64(len(payload))
			} else if total == 1 {
				cSize = int64(len(payload))
			} else if existingChunks, ok := m.chunks[token]; ok {
				for sIdx, cData := range existingChunks {
					if sIdx < total && len(cData) > 0 {
						cSize = int64(len(cData))
						break
					}
				}
			}

			if cSize > 0 {
				m.chunkSizes[token] = cSize
				// Canonical chunk size just discovered; flush any previously buffered chunks (e.g. remainder chunk seq == total)
				for sIdx, cData := range tokenChunks {
					if len(cData) > 0 {
						remOffset := int64(sIdx-1) * cSize
						_, _ = diskF.WriteAt(cData, remOffset)
						tokenChunks[sIdx] = nil
					}
				}
			}
		}

		if cSize > 0 {
			offset := int64(seq-1) * cSize
			_, _ = diskF.WriteAt(payload, offset)
			tokenChunks[seq] = nil
		} else {
			// Canonical chunk size not yet known (remainder chunk arrived first for total > 1).
			// Buffer chunk in memory until canonical size is determined.
			chunkCopy := make([]byte, len(payload))
			copy(chunkCopy, payload)
			tokenChunks[seq] = chunkCopy
		}

		session.ReceivedChunks = len(tokenChunks)
		session.TotalBytes += int64(len(payload))
		session.LastActive = time.Now().UTC()
		if len(session.Chunks) < 32 {
			session.Chunks[seq] = common.ChunkRecord{
				Sequence:   seq,
				Total:      total,
				Length:     uint32(len(payload)),
				ReceivedAt: time.Now().UTC(),
			}
		}
	} else {
		tokenChunks := m.chunks[token]
		if _, duplicate := tokenChunks[seq]; !duplicate {
			chunkCopy := make([]byte, len(payload))
			copy(chunkCopy, payload)
			tokenChunks[seq] = chunkCopy

			session.ReceivedChunks = len(tokenChunks)
			session.TotalBytes += int64(len(payload))
			session.LastActive = time.Now().UTC()
			if len(session.Chunks) < 32 {
				session.Chunks[seq] = common.ChunkRecord{
					Sequence:   seq,
					Total:      total,
					Length:     uint32(len(payload)),
					ReceivedAt: time.Now().UTC(),
				}
			}
		}
	}

	obs := m.observer

	// Check if all chunks have arrived
	if session.ReceivedChunks >= int(total) {
		var reasmErr error
		if isSpilled && diskF != nil {
			// Ensure any buffered chunk (e.g. remainder chunk) is flushed to disk
			if cSize := m.chunkSizes[token]; cSize > 0 {
				if tokenChunks, ok := m.chunks[token]; ok {
					for sIdx, cData := range tokenChunks {
						if len(cData) > 0 {
							offset := int64(sIdx-1) * cSize
							_, _ = diskF.WriteAt(cData, offset)
							tokenChunks[sIdx] = nil
						}
					}
				}
			}
			_ = diskF.Sync()
			_, _ = diskF.Seek(0, 0)
			hasher := sha256.New()
			buf := make([]byte, 64*1024)
			_, _ = io.CopyBuffer(hasher, diskF, buf)
			computedHash := hex.EncodeToString(hasher.Sum(nil))

			if session.ExpectedHash != "" && !strings.EqualFold(computedHash, session.ExpectedHash) {
				session.PayloadHash = computedHash
				session.Status = "corrupted"
				session.Verdict = common.VerdictBlockedDLP
				reasmErr = ErrHashMismatch
			} else {
				session.PayloadHash = computedHash
				session.Status = "completed"
				session.Verdict = common.VerdictSuccess
			}
			tmpPath := diskF.Name()
			_ = diskF.Close()
			delete(m.diskFiles, token)
			delete(m.chunks, token)
			delete(m.chunkSizes, token)
			if m.diskPaths != nil {
				m.diskPaths[token] = tmpPath
			}
			m.completedTokens = append(m.completedTokens, token)
			if len(m.completedTokens) > 2000 {
				oldTok := m.completedTokens[0]
				m.completedTokens = m.completedTokens[1:]
				delete(m.chunks, oldTok)
				if m.diskPaths != nil {
					if oldPath, ok := m.diskPaths[oldTok]; ok {
						_ = os.Remove(oldPath)
						delete(m.diskPaths, oldTok)
					}
				}
			}
		} else {
			tokenChunks := m.chunks[token]
			res := ReassembleChunks(tokenChunks, total, session.ExpectedHash)
			session.PayloadHash = res.ComputedHash
			session.Status = res.Status
			session.Verdict = res.Verdict

			m.completedTokens = append(m.completedTokens, token)
			if len(m.completedTokens) > 2000 {
				oldTok := m.completedTokens[0]
				m.completedTokens = m.completedTokens[1:]
				delete(m.chunks, oldTok)
				if m.diskPaths != nil {
					if oldPath, ok := m.diskPaths[oldTok]; ok {
						_ = os.Remove(oldPath)
						delete(m.diskPaths, oldTok)
					}
				}
			}

			if res.Error != nil {
				reasmErr = res.Error
			}
		}

		if reasmErr != nil {
			retSession := CloneSession(session)
			m.mu.Unlock()
			if obs != nil {
				obs.OnChunkReceived(retSession, seq, total, int64(len(payload)))
				obs.OnSessionCompleted(retSession)
			}
			return retSession, true, reasmErr
		}

		retSession := CloneSession(session)
		m.mu.Unlock()

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

		return retSession, true, nil
	}

	session.Status = "active"
	session.Verdict = common.VerdictBlockedL7
	retSession := CloneSession(session)
	m.mu.Unlock()

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

	return retSession, false, nil
}

// GetSessionPayload reconstructs the full contiguous payload for a session.
func (m *MemoryLedger) GetSessionPayload(token common.SessionToken) ([]byte, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	if m.closed {
		return nil, ErrLedgerClosed
	}

	session, exists := m.sessions[token]
	if !exists {
		return nil, ErrSessionNotFound
	}

	if session.Status == "corrupted" || session.Verdict == common.VerdictBlockedDLP {
		return nil, ErrHashMismatch
	}

	if m.diskPaths != nil {
		if p, ok := m.diskPaths[token]; ok && p != "" {
			return os.ReadFile(p)
		}
	}

	if m.diskFiles != nil {
		if df, ok := m.diskFiles[token]; ok && df != nil {
			_ = df.Sync()
			_, _ = df.Seek(0, 0)
			return io.ReadAll(df)
		}
	}

	chunks, exists := m.chunks[token]
	if !exists || len(chunks) == 0 {
		return nil, ErrSessionIncomplete
	}

	total := uint32(session.TotalChunks)
	if total == 0 {
		return nil, fmt.Errorf("%w: total chunks is 0", ErrSessionIncomplete)
	}
	if len(chunks) < int(total) {
		return nil, fmt.Errorf("%w: received %d of %d chunks", ErrSessionIncomplete, len(chunks), total)
	}

	var fullPayload []byte
	for s := uint32(1); s <= total; s++ {
		chunkData, found := chunks[s]
		if !found {
			return nil, fmt.Errorf("%w: missing chunk sequence %d", ErrSessionIncomplete, s)
		}
		fullPayload = append(fullPayload, chunkData...)
	}

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

	return fullPayload, nil
}

// RecordCredential stores harvested credentials in memory.
func (m *MemoryLedger) RecordCredential(cred common.CapturedCredential) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.closed {
		return ErrLedgerClosed
	}

	m.credentials = append(m.credentials, cred)
	return nil
}

// GetCredentials returns a snapshot of all captured credentials.
func (m *MemoryLedger) GetCredentials() ([]common.CapturedCredential, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	if m.closed {
		return nil, ErrLedgerClosed
	}

	copied := make([]common.CapturedCredential, len(m.credentials))
	copy(copied, m.credentials)
	return copied, nil
}

// ExportJSON serializes the entire ledger state into JSON.
func (m *MemoryLedger) ExportJSON() ([]byte, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	if m.closed {
		return nil, ErrLedgerClosed
	}

	probes := make([]common.ProbeResult, 0, len(m.probes))
	for _, p := range m.probes {
		probes = append(probes, p)
	}

	sessions := make([]*common.Session, 0, len(m.sessions))
	for _, s := range m.sessions {
		sessions = append(sessions, CloneSession(s))
	}

	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(m.credentials),
		Probes:           probes,
		Sessions:         sessions,
		Credentials:      m.credentials,
	}

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

// Reset clears all recorded probes, sessions, chunks, credentials, and disk buffer files.
func (m *MemoryLedger) Reset() error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.closed {
		return ErrLedgerClosed
	}

	for tok, f := range m.diskFiles {
		_ = f.Close()
		_ = os.Remove(f.Name())
		delete(m.diskFiles, tok)
	}
	for tok, p := range m.diskPaths {
		_ = os.Remove(p)
		delete(m.diskPaths, tok)
	}

	m.probes = make(map[string]common.ProbeResult)
	m.sessions = make(map[common.SessionToken]*common.Session)
	m.chunks = make(map[common.SessionToken]map[uint32][]byte)
	m.chunkSizes = make(map[common.SessionToken]int64)
	m.completedTokens = make([]common.SessionToken, 0, 64)
	m.credentials = make([]common.CapturedCredential, 0)
	return nil
}

// Close marks the ledger as closed and purges stored chunks.
func (m *MemoryLedger) Close() error {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.closed = true
	m.chunks = make(map[common.SessionToken]map[uint32][]byte)
	m.chunkSizes = make(map[common.SessionToken]int64)
	for tok, f := range m.diskFiles {
		_ = f.Close()
		_ = os.Remove(f.Name())
		delete(m.diskFiles, tok)
	}
	for tok, p := range m.diskPaths {
		_ = os.Remove(p)
		delete(m.diskPaths, tok)
	}
	return nil
}