LOGAR/tests/test_hub_storage_deep.py

"""
Comprehensive unit tests for LOGAR StorageEngine, StorageCryptoManager, and Multi-Backend Migrations.
"""

import os
import sys
import json
import shutil
import tempfile
import sqlite3
import unittest
from unittest.mock import MagicMock, patch

from src.hub.storage_engine import StorageEngine, StorageCryptoManager


class TestStorageEngineDeep(unittest.TestCase):
    def setUp(self):
        self.temp_dir = tempfile.mkdtemp()
        self.db_path = os.path.join(self.temp_dir, "test_logar.db")

    def tearDown(self):
        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_storage_crypto_manager_keypair_and_encryption(self):
        """Tests keypair generation, field encryption, decryption, and error resilience."""
        priv, pub, fp = StorageCryptoManager.generate_server_keypair("Test Server Key")
        self.assertIn("BEGIN PGP PRIVATE KEY", priv)
        self.assertIn("BEGIN PGP PUBLIC KEY", pub)
        self.assertTrue(len(fp) > 0)

        server_state = {
            "config": {
                "encrypt_at_rest": True,
                "public_key": pub,
                "private_key": priv
            }
        }

        # Test empty or non-string inputs
        self.assertEqual(StorageCryptoManager.encrypt_field_for_storage("", server_state=server_state, is_unlocked=True), "")
        self.assertEqual(StorageCryptoManager.encrypt_field_for_storage(None, server_state=server_state, is_unlocked=True), None)
        self.assertEqual(StorageCryptoManager.decrypt_field_from_storage("", server_state=server_state), "")
        self.assertEqual(StorageCryptoManager.decrypt_field_from_storage(None, server_state=server_state), None)

        # Test non-encrypted plain string decrypt passthrough
        plain_msg = "Security Audit Message 2026"
        self.assertEqual(StorageCryptoManager.decrypt_field_from_storage(plain_msg, server_state=server_state), plain_msg)

        # Encrypt when unlocked
        encrypted = StorageCryptoManager.encrypt_field_for_storage(plain_msg, server_state=server_state, is_unlocked=True)
        self.assertTrue(encrypted.startswith("-----BEGIN PGP MESSAGE-----"))

        # Encrypt idempotence (should not double-encrypt)
        double_enc = StorageCryptoManager.encrypt_field_for_storage(encrypted, server_state=server_state, is_unlocked=True)
        self.assertEqual(double_enc, encrypted)

        # Decrypt
        decrypted = StorageCryptoManager.decrypt_field_from_storage(encrypted, server_state=server_state)
        self.assertEqual(decrypted, plain_msg)

        # Decrypt with state missing private_key_obj (reloads from blob)
        clean_state = {"config": {"private_key": priv}}
        decrypted_from_clean = StorageCryptoManager.decrypt_field_from_storage(encrypted, server_state=clean_state)
        self.assertEqual(decrypted_from_clean, plain_msg)

        # Encrypt when locked / encrypt_at_rest is false
        unenc_state = {"config": {"encrypt_at_rest": False, "public_key": pub}}
        res_plain = StorageCryptoManager.encrypt_field_for_storage(plain_msg, server_state=unenc_state, is_unlocked=False)
        self.assertEqual(res_plain, plain_msg)

    def test_storage_engine_sqlite_init_and_telemetry(self):
        """Tests SQLite database initialization, WAL settings, column migrations, and telemetry."""
        db = StorageEngine.init_db(self.db_path, enrollment_secret="sec-test-123", max_seats=50)
        self.assertEqual(db, self.db_path)
        self.assertTrue(os.path.exists(db))

        # Re-running init_db executes migrations smoothly
        db_reinit = StorageEngine.init_db(self.db_path, enrollment_secret="sec-test-123", max_seats=50)
        self.assertEqual(db_reinit, self.db_path)

        # Check telemetry info
        state = {"config": {"db_path": self.db_path, "db_type": "sqlite"}}
        info = StorageEngine.get_storage_engine_info(self.db_path, server_state=state, is_unlocked=True)
        self.assertEqual(info["engine"], "sqlite")
        self.assertFalse(info["is_enterprise"])
        self.assertIn("active_issues", info["table_stats"])
        self.assertIn("clients", info["table_stats"])

    def test_storage_engine_postgres_mocked(self):
        """Tests Postgres schema initialization and SQLite-to-Postgres data migration with mocked drivers."""
        pg_cfg = {
            "host": "127.0.0.1",
            "port": 5432,
            "database": "logar_test",
            "user": "logar_user",
            "password": "secret_password",
            "sslmode": "disable"
        }

        # Initialize SQLite DB with test records
        StorageEngine.init_db(self.db_path, enrollment_secret="sec-migrate", max_seats=10)
        conn = sqlite3.connect(self.db_path)
        conn.execute("""
            INSERT INTO active_issues (fingerprint, site_name, server, signature, severity, message, os_type, first_seen, last_seen, run_count, status, last_run_id)
            VALUES ('fp1', 'SiteA', 'srv1', 'Sig1', 'HIGH', 'Test warning message', 'linux', '2026-01-01', '2026-01-02', 1, 'active', 'run1')
        """)
        conn.commit()
        conn.close()

        # Mock psycopg2
        mock_psycopg2 = MagicMock()
        mock_conn = MagicMock()
        mock_cur = MagicMock()
        mock_conn.cursor.return_value.__enter__.return_value = mock_cur
        mock_psycopg2.connect.return_value = mock_conn

        with patch.dict(sys.modules, {"psycopg2": mock_psycopg2}):
            # Test Postgres DB init
            ok, msg = StorageEngine.init_postgres_db(pg_cfg)
            self.assertTrue(ok)
            self.assertIn("successfully initialized", msg)

            # Test Migration
            encrypt_mock = lambda x: f"ENC({x})"
            ok, msg, counts = StorageEngine.migrate_sqlite_to_postgres(self.db_path, pg_cfg, encrypt_fn=encrypt_mock)
            self.assertTrue(ok)
            self.assertIn("active_issues", counts)
            self.assertEqual(counts["active_issues"], 1)

            # Test Postgres info retrieval
            mock_cur.fetchone.return_value = (10,)
            state = {"config": {"db_type": "postgres", "postgres": pg_cfg}}
            info = StorageEngine.get_storage_engine_info(None, server_state=state, is_unlocked=True)
            self.assertEqual(info["engine"], "postgres")
            self.assertTrue(info["is_enterprise"])
            self.assertEqual(info["table_stats"]["active_issues"], 10)


if __name__ == "__main__":
    unittest.main()