diff --git a/node/beacon_signature_verification.py b/node/beacon_signature_verification.py new file mode 100644 index 0000000..18ee63f --- /dev/null +++ b/node/beacon_signature_verification.py @@ -0,0 +1,216 @@ +""" +Beacon Atlas Signature Verification Module +Implements Ed25519 signature verification for /relay/ping endpoint using TOFU stored public keys. + +This module provides the core verification logic that should be integrated into the +Beacon Atlas Flask application's /relay/ping endpoint. + +Usage: + from beacon_signature_verification import verify_relay_ping_signature + + # In your /relay/ping endpoint handler + if agent_has_stored_pubkey(agent_id): + if not verify_relay_ping_signature(agent_id, request_data, signature): + return jsonify({"error": "Invalid signature"}), 401 +""" + +import base64 +import json +import sqlite3 +from typing import Optional, Dict, Any + +# Import existing crypto utilities from beacon_skill +try: + from beacon_skill.crypto import verify_ed25519_signature + HAVE_BEACON_SKILL = True +except ImportError: + HAVE_BEACON_SKILL = False + +# Fallback Ed25519 verification using pynacl (same as TOFU implementation) +try: + from nacl.signing import VerifyKey + from nacl.exceptions import BadSignatureError + HAVE_NACL = True +except ImportError: + HAVE_NACL = False + + +def get_db() -> sqlite3.Connection: + """ + Get database connection to beacon_atlas.db. + This should be replaced with the actual database connection method used by Beacon Atlas. + """ + db_path = "/root/rustchain/node/beacon_atlas.db" + return sqlite3.connect(db_path) + + +def tofu_get_key_info(conn: sqlite3.Connection, agent_id: str) -> Optional[Dict[str, Any]]: + """ + Retrieve TOFU key information for an agent from the relay_agents table. + + Args: + conn: Database connection + agent_id: Agent identifier + + Returns: + Dictionary containing key info or None if not found + """ + try: + cursor = conn.execute( + "SELECT pubkey_hex, revoked FROM relay_agents WHERE agent_id = ?", + (agent_id,) + ) + row = cursor.fetchone() + if row: + return { + "pubkey_hex": row[0], + "revoked": bool(row[1]) + } + return None + except sqlite3.OperationalError: + # Table might not exist yet + return None + + +def verify_ed25519_fallback(pubkey_hex: str, message: bytes, signature_b64: str) -> bool: + """ + Fallback Ed25519 signature verification using pynacl. + + Args: + pubkey_hex: Public key in hexadecimal format + message: Message bytes to verify + signature_b64: Base64-encoded signature + + Returns: + True if signature is valid, False otherwise + """ + if not HAVE_NACL: + return False + + try: + # Convert hex pubkey to bytes + pubkey_bytes = bytes.fromhex(pubkey_hex) + verify_key = VerifyKey(pubkey_bytes) + + # Decode base64 signature + signature_bytes = base64.b64decode(signature_b64) + + # Verify signature + verify_key.verify(message, signature_bytes) + return True + except (BadSignatureError, ValueError, TypeError): + return False + + +def verify_relay_ping_signature(agent_id: str, payload: Dict[str, Any], signature: str) -> bool: + """ + Verify Ed25519 signature for /relay/ping endpoint using TOFU stored public key. + + This function should be called by the Beacon Atlas /relay/ping endpoint handler + when signature verification is required. + + Args: + agent_id: Relay agent ID from the ping request + payload: Ping payload data (excluding signature field) + signature: Base64-encoded Ed25519 signature of the payload + + Returns: + bool: True if signature is valid and agent has a valid stored pubkey, + False otherwise (including cases where agent has no stored pubkey) + """ + # Get database connection + try: + conn = get_db() + except Exception: + return False + + try: + # Retrieve stored public key for this agent + key_info = tofu_get_key_info(conn, agent_id) + + # If no key info or key is revoked, verification fails + if not key_info or key_info.get('revoked', False): + return False + + pubkey_hex = key_info['pubkey_hex'] + if not pubkey_hex: + return False + + # Prepare message for verification (JSON payload without signature) + # Ensure consistent JSON serialization + message_dict = payload.copy() + # Remove signature field if present to avoid circular dependency + if 'signature' in message_dict: + del message_dict['signature'] + + message_json = json.dumps(message_dict, sort_keys=True, separators=(',', ':')) + message_bytes = message_json.encode('utf-8') + + # Verify signature using available crypto library + if HAVE_BEACON_SKILL: + try: + return verify_ed25519_signature(pubkey_hex, message_bytes, signature) + except Exception: + # Fall back to pynacl if beacon_skill fails + pass + + # Use fallback verification + return verify_ed25519_fallback(pubkey_hex, message_bytes, signature) + + finally: + conn.close() + + +def agent_has_stored_pubkey(agent_id: str) -> bool: + """ + Check if an agent has a stored public key in the TOFU database. + This function helps determine whether signature verification should be enforced. + + Args: + agent_id: Agent identifier + + Returns: + True if agent has a stored pubkey, False otherwise + """ + try: + conn = get_db() + key_info = tofu_get_key_info(conn, agent_id) + conn.close() + return key_info is not None and key_info.get('pubkey_hex') is not None + except Exception: + return False + + +# Integration example for Beacon Atlas /relay/ping endpoint +def integrate_with_beacon_atlas_example(): + """ + Example of how to integrate this module with the Beacon Atlas Flask application. + + In your Beacon Atlas application (likely in a file like beacon_atlas.py or similar): + + ```python + from beacon_signature_verification import verify_relay_ping_signature, agent_has_stored_pubkey + from flask import request, jsonify + + @app.route('/relay/ping', methods=['POST']) + def relay_ping(): + data = request.get_json() + agent_id = data.get('agent_id') + + # Backward compatibility: only enforce signature verification for agents with stored pubkeys + if agent_has_stored_pubkey(agent_id): + signature = data.get('signature') + if not signature: + return jsonify({"error": "Signature required"}), 401 + + # Verify signature + if not verify_relay_ping_signature(agent_id, data, signature): + return jsonify({"error": "Invalid signature"}), 401 + + # Process ping as usual + # ... rest of ping handling logic ... + + return jsonify({"status": "ok"}) + ``` + """ + pass \ No newline at end of file diff --git a/tests/test_beacon_signature_verification.py b/tests/test_beacon_signature_verification.py new file mode 100644 index 0000000..dbd23f4 --- /dev/null +++ b/tests/test_beacon_signature_verification.py @@ -0,0 +1,226 @@ +""" +Test suite for beacon_signature_verification module. +""" + +import unittest +import tempfile +import os +import sqlite3 +import json +import base64 +from unittest.mock import patch, MagicMock + +# Add the node directory to Python path +import sys +sys.path.insert(0, 'node') + +from beacon_signature_verification import ( + verify_relay_ping_signature, + agent_has_stored_pubkey, + tofu_get_key_info, + verify_ed25519_fallback +) + + +class TestBeaconSignatureVerification(unittest.TestCase): + """Test cases for beacon signature verification.""" + + def setUp(self): + """Set up test database and temporary files.""" + self.temp_dir = tempfile.mkdtemp() + self.db_path = os.path.join(self.temp_dir, "test_beacon_atlas.db") + + # Create test database with relay_agents table + conn = sqlite3.connect(self.db_path) + conn.execute(""" + CREATE TABLE IF NOT EXISTS relay_agents ( + agent_id TEXT PRIMARY KEY, + pubkey_hex TEXT, + revoked INTEGER DEFAULT 0, + last_seen REAL, + created_at REAL + ) + """) + conn.commit() + conn.close() + + # Mock the get_db function to use our test database + self.original_get_db = None + + def tearDown(self): + """Clean up test files.""" + if os.path.exists(self.db_path): + os.remove(self.db_path) + os.rmdir(self.temp_dir) + + def mock_get_db(self): + """Mock database connection function.""" + return sqlite3.connect(self.db_path) + + def test_tofu_get_key_info_found(self): + """Test retrieving key info for existing agent.""" + # Insert test data + conn = sqlite3.connect(self.db_path) + conn.execute( + "INSERT INTO relay_agents (agent_id, pubkey_hex, revoked, created_at) VALUES (?, ?, ?, ?)", + ("test_agent_1", "a1b2c3d4e5f6..." * 4, 0, 1234567890) + ) + conn.commit() + conn.close() + + # Test retrieval + with patch('beacon_signature_verification.get_db', self.mock_get_db): + key_info = tofu_get_key_info(None, "test_agent_1") + self.assertIsNotNone(key_info) + self.assertEqual(key_info["pubkey_hex"], "a1b2c3d4e5f6..." * 4) + self.assertFalse(key_info["revoked"]) + + def test_tofu_get_key_info_not_found(self): + """Test retrieving key info for non-existent agent.""" + with patch('beacon_signature_verification.get_db', self.mock_get_db): + key_info = tofu_get_key_info(None, "nonexistent_agent") + self.assertIsNone(key_info) + + def test_tofu_get_key_info_revoked(self): + """Test retrieving key info for revoked agent.""" + # Insert test data + conn = sqlite3.connect(self.db_path) + conn.execute( + "INSERT INTO relay_agents (agent_id, pubkey_hex, revoked, created_at) VALUES (?, ?, ?, ?)", + ("revoked_agent", "b2c3d4e5f6g7..." * 4, 1, 1234567890) + ) + conn.commit() + conn.close() + + # Test retrieval + with patch('beacon_signature_verification.get_db', self.mock_get_db): + key_info = tofu_get_key_info(None, "revoked_agent") + self.assertIsNotNone(key_info) + self.assertTrue(key_info["revoked"]) + + def test_agent_has_stored_pubkey_true(self): + """Test agent has stored pubkey returns True.""" + # Insert test data + conn = sqlite3.connect(self.db_path) + conn.execute( + "INSERT INTO relay_agents (agent_id, pubkey_hex, created_at) VALUES (?, ?, ?)", + ("test_agent_2", "c3d4e5f6g7h8..." * 4, 1234567890) + ) + conn.commit() + conn.close() + + with patch('beacon_signature_verification.get_db', self.mock_get_db): + result = agent_has_stored_pubkey("test_agent_2") + self.assertTrue(result) + + def test_agent_has_stored_pubkey_false(self): + """Test agent has stored pubkey returns False for non-existent agent.""" + with patch('beacon_signature_verification.get_db', self.mock_get_db): + result = agent_has_stored_pubkey("nonexistent_agent") + self.assertFalse(result) + + def test_verify_ed25519_fallback_no_nacl(self): + """Test fallback verification when pynacl is not available.""" + with patch('beacon_signature_verification.HAVE_NACL', False): + result = verify_ed25519_fallback("a1b2c3d4e5f6...", b"test message", "signature") + self.assertFalse(result) + + @patch('beacon_signature_verification.HAVE_NACL', True) + @patch('beacon_signature_verification.VerifyKey') + def test_verify_ed25519_fallback_valid_signature(self, mock_verify_key): + """Test fallback verification with valid signature.""" + mock_verify_instance = MagicMock() + mock_verify_key.return_value = mock_verify_instance + mock_verify_instance.verify.return_value = None # No exception means valid + + result = verify_ed25519_fallback("a1b2c3d4e5f6..." * 8, b"test message", "dGVzdCBzaWduYXR1cmU=") + self.assertTrue(result) + mock_verify_instance.verify.assert_called_once() + + @patch('beacon_signature_verification.HAVE_NACL', True) + @patch('beacon_signature_verification.VerifyKey') + def test_verify_ed25519_fallback_invalid_signature(self, mock_verify_key): + """Test fallback verification with invalid signature.""" + from nacl.exceptions import BadSignatureError + mock_verify_instance = MagicMock() + mock_verify_key.return_value = mock_verify_instance + mock_verify_instance.verify.side_effect = BadSignatureException() + + result = verify_ed25519_fallback("a1b2c3d4e5f6..." * 8, b"test message", "invalid_signature") + self.assertFalse(result) + + def test_verify_relay_ping_signature_no_db_connection(self): + """Test verification fails when database connection fails.""" + with patch('beacon_signature_verification.get_db', side_effect=Exception("DB error")): + result = verify_relay_ping_signature("test_agent", {}, "signature") + self.assertFalse(result) + + def test_verify_relay_ping_signature_no_key_info(self): + """Test verification fails when no key info exists.""" + with patch('beacon_signature_verification.get_db', self.mock_get_db): + result = verify_relay_ping_signature("nonexistent_agent", {}, "signature") + self.assertFalse(result) + + def test_verify_relay_ping_signature_revoked_key(self): + """Test verification fails when key is revoked.""" + # Insert revoked key + conn = sqlite3.connect(self.db_path) + conn.execute( + "INSERT INTO relay_agents (agent_id, pubkey_hex, revoked, created_at) VALUES (?, ?, ?, ?)", + ("revoked_test", "d4e5f6g7h8i9..." * 4, 1, 1234567890) + ) + conn.commit() + conn.close() + + with patch('beacon_signature_verification.get_db', self.mock_get_db): + result = verify_relay_ping_signature("revoked_test", {}, "signature") + self.assertFalse(result) + + @patch('beacon_signature_verification.HAVE_BEACON_SKILL', False) + @patch('beacon_signature_verification.HAVE_NACL', True) + @patch('beacon_signature_verification.VerifyKey') + def test_verify_relay_ping_signature_fallback_success(self, mock_verify_key): + """Test full verification flow using fallback crypto.""" + # Insert test key + conn = sqlite3.connect(self.db_path) + conn.execute( + "INSERT INTO relay_agents (agent_id, pubkey_hex, created_at) VALUES (?, ?, ?)", + ("test_agent_3", "e5f6g7h8i9j0..." * 4, 1234567890) + ) + conn.commit() + conn.close() + + # Mock successful verification + mock_verify_instance = MagicMock() + mock_verify_key.return_value = mock_verify_instance + mock_verify_instance.verify.return_value = None + + # Test payload + test_payload = { + "agent_id": "test_agent_3", + "timestamp": 1234567890, + "status": "alive" + } + + with patch('beacon_signature_verification.get_db', self.mock_get_db): + result = verify_relay_ping_signature("test_agent_3", test_payload, "valid_signature") + self.assertTrue(result) + + def test_verify_relay_ping_signature_payload_serialization(self): + """Test that payload serialization is consistent.""" + # This test ensures that the JSON serialization produces consistent results + # which is crucial for signature verification + + payload1 = {"a": 1, "b": 2, "c": 3} + payload2 = {"c": 3, "b": 2, "a": 1} # Same content, different order + + # Both should produce the same JSON string due to sort_keys=True + json1 = json.dumps(payload1, sort_keys=True, separators=(',', ':')) + json2 = json.dumps(payload2, sort_keys=True, separators=(',', ':')) + + self.assertEqual(json1, json2) + self.assertEqual(json1, '{"a":1,"b":2,"c":3}') + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file