File size: 9,347 Bytes
631d977
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
"""
Ledger Store Module
====================
Append-only ledger storage with merkle root computation.
Provides immutable audit trail for agent actions.
"""

import json
import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional
from dataclasses import dataclass, field, asdict

from .merkle import compute_merkle_root, hash_leaf


@dataclass
class LedgerEntry:
    """Represents a single entry in the ledger."""
    id: str
    timestamp: str
    action_type: str
    agent: str
    input_hash: str
    output_hash: str
    data: dict[str, Any]
    previous_hash: str
    entry_hash: str = ""
    
    def __post_init__(self):
        if not self.entry_hash:
            self.entry_hash = self._compute_hash()
    
    def _compute_hash(self) -> str:
        """Compute hash of this entry."""
        content = f"{self.id}:{self.timestamp}:{self.action_type}:{self.agent}:{self.previous_hash}"
        return hash_leaf(content)
    
    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class LedgerBlock:
    """Represents a block of ledger entries."""
    block_id: int
    timestamp: str
    entries: list[LedgerEntry]
    previous_block_hash: str
    merkle_root: str = ""
    block_hash: str = ""
    signature: str = ""  # Will be filled by ECDSA signing
    
    def __post_init__(self):
        if not self.merkle_root:
            self.merkle_root = self._compute_merkle_root()
        if not self.block_hash:
            self.block_hash = self._compute_hash()
    
    def _compute_merkle_root(self) -> str:
        """Compute merkle root from entry hashes."""
        if not self.entries:
            return hash_leaf("empty_block")
        entry_hashes = [entry.entry_hash for entry in self.entries]
        return compute_merkle_root(entry_hashes)
    
    def _compute_hash(self) -> str:
        """Compute hash of this block."""
        content = f"{self.block_id}:{self.timestamp}:{self.merkle_root}:{self.previous_block_hash}"
        return hash_leaf(content)
    
    def to_dict(self) -> dict[str, Any]:
        return {
            "block_id": self.block_id,
            "timestamp": self.timestamp,
            "entries": [e.to_dict() for e in self.entries],
            "previous_block_hash": self.previous_block_hash,
            "merkle_root": self.merkle_root,
            "block_hash": self.block_hash,
            "signature": self.signature
        }


class LedgerStore:
    """
    Append-only ledger store.
    
    Provides:
    - Immutable storage of agent actions
    - Block-based organization with merkle roots
    - JSON persistence
    - Verification of chain integrity
    """
    
    GENESIS_HASH = "0" * 64  # Genesis block previous hash
    
    def __init__(self, storage_path: Optional[str] = None):
        """
        Initialize the ledger store.
        
        Args:
            storage_path: Path to store ledger JSON file. If None, uses in-memory only.
        """
        self.storage_path = Path(storage_path) if storage_path else None
        self.blocks: list[LedgerBlock] = []
        self.pending_entries: list[LedgerEntry] = []
        self._entry_counter = 0
        
        # Load existing ledger if available
        if self.storage_path and self.storage_path.exists():
            self._load()
    
    def add_entry(
        self,
        action_type: str,
        agent: str,
        input_data: dict[str, Any],
        output_data: dict[str, Any]
    ) -> LedgerEntry:
        """
        Add a new entry to the pending entries.
        
        Args:
            action_type: Type of action (e.g., "search", "summarize")
            agent: Name of the agent performing the action
            input_data: Input data for the action
            output_data: Output data from the action
            
        Returns:
            The created LedgerEntry
        """
        self._entry_counter += 1
        
        # Get previous hash (from last pending entry or last block)
        if self.pending_entries:
            previous_hash = self.pending_entries[-1].entry_hash
        elif self.blocks:
            previous_hash = self.blocks[-1].block_hash
        else:
            previous_hash = self.GENESIS_HASH
        
        entry = LedgerEntry(
            id=f"entry_{self._entry_counter}",
            timestamp=datetime.now(timezone.utc).isoformat(),
            action_type=action_type,
            agent=agent,
            input_hash=hash_leaf(json.dumps(input_data, sort_keys=True)),
            output_hash=hash_leaf(json.dumps(output_data, sort_keys=True)),
            data={"input": input_data, "output": output_data},
            previous_hash=previous_hash
        )
        
        self.pending_entries.append(entry)
        return entry
    
    def create_block(self, signer=None) -> LedgerBlock:
        """
        Create a new block from pending entries.
        
        Args:
            signer: Optional signer object with sign() method for ECDSA signing
            
        Returns:
            The created LedgerBlock
        """
        if not self.pending_entries:
            raise ValueError("No pending entries to create block")
        
        # Get previous block hash
        if self.blocks:
            previous_block_hash = self.blocks[-1].block_hash
        else:
            previous_block_hash = self.GENESIS_HASH
        
        block = LedgerBlock(
            block_id=len(self.blocks) + 1,
            timestamp=datetime.now(timezone.utc).isoformat(),
            entries=self.pending_entries.copy(),
            previous_block_hash=previous_block_hash
        )
        
        # Sign the block if signer is provided
        if signer:
            block.signature = signer.sign(block.block_hash)
        
        self.blocks.append(block)
        self.pending_entries = []
        
        # Persist to storage
        if self.storage_path:
            self._save()
        
        return block
    
    def get_block(self, block_id: int) -> Optional[LedgerBlock]:
        """Get a block by ID."""
        for block in self.blocks:
            if block.block_id == block_id:
                return block
        return None
    
    def get_latest_block(self) -> Optional[LedgerBlock]:
        """Get the most recent block."""
        return self.blocks[-1] if self.blocks else None
    
    def verify_chain(self) -> tuple[bool, str]:
        """
        Verify the integrity of the entire blockchain.
        
        Returns:
            Tuple of (is_valid, message)
        """
        if not self.blocks:
            return True, "Empty chain is valid"
        
        # Check genesis block
        if self.blocks[0].previous_block_hash != self.GENESIS_HASH:
            return False, "Genesis block has invalid previous hash"
        
        # Check each block's link
        for i in range(1, len(self.blocks)):
            current = self.blocks[i]
            previous = self.blocks[i - 1]
            
            if current.previous_block_hash != previous.block_hash:
                return False, f"Block {current.block_id} has invalid previous hash"
            
            # Verify merkle root
            computed_root = current._compute_merkle_root()
            if computed_root != current.merkle_root:
                return False, f"Block {current.block_id} has invalid merkle root"
        
        return True, "Chain is valid"
    
    def get_full_ledger(self) -> dict[str, Any]:
        """Get the complete ledger as a dictionary."""
        return {
            "blocks": [block.to_dict() for block in self.blocks],
            "pending_entries": [entry.to_dict() for entry in self.pending_entries],
            "total_blocks": len(self.blocks),
            "total_entries": sum(len(b.entries) for b in self.blocks) + len(self.pending_entries)
        }
    
    def _save(self):
        """Save ledger to JSON file."""
        if not self.storage_path:
            return
        
        # Ensure directory exists
        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
        
        data = self.get_full_ledger()
        with open(self.storage_path, 'w') as f:
            json.dump(data, f, indent=2)
    
    def _load(self):
        """Load ledger from JSON file."""
        if not self.storage_path or not self.storage_path.exists():
            return
        
        with open(self.storage_path, 'r') as f:
            data = json.load(f)
        
        # Reconstruct blocks
        self.blocks = []
        for block_data in data.get("blocks", []):
            entries = [
                LedgerEntry(**entry_data)
                for entry_data in block_data.get("entries", [])
            ]
            block = LedgerBlock(
                block_id=block_data["block_id"],
                timestamp=block_data["timestamp"],
                entries=entries,
                previous_block_hash=block_data["previous_block_hash"],
                merkle_root=block_data["merkle_root"],
                block_hash=block_data["block_hash"],
                signature=block_data.get("signature", "")
            )
            self.blocks.append(block)
        
        # Update entry counter
        total_entries = sum(len(b.entries) for b in self.blocks)
        self._entry_counter = total_entries