import asyncio import os import logging import uuid from datetime import datetime, timedelta from motor.motor_asyncio import AsyncIOMotorClient from pymongo import IndexModel, ASCENDING, DESCENDING, TEXT from dotenv import load_dotenv # Load environment variables load_dotenv() # Configure Logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) MONGODB_URL = os.getenv("MONGODB_URL", "mongodb://localhost:27017") DATABASE_NAME = "nadraguard_db" async def setup_database(): logger.info("Initializing NadraGuard Database Setup...") client = AsyncIOMotorClient(MONGODB_URL) db = client[DATABASE_NAME] # --- 1. USERS COLLECTION --- logger.info("Setting up 'users' collection...") user_validator = { "$jsonSchema": { "bsonType": "object", "required": ["user_id", "phone", "language_preference", "is_active"], "properties": { "user_id": {"bsonType": "string"}, "phone": {"bsonType": "string", "pattern": "^03\\d{2}-\\d{7}$"}, "language_preference": {"enum": ["ur", "roman_ur", "en"]}, "family_circle": { "bsonType": "array", "items": { "bsonType": "object", "properties": { "name": {"bsonType": "string"}, "phone": {"bsonType": "string"}, "language": {"bsonType": "string"}, "relationship": {"bsonType": "string"} } } }, "wallet_accounts": { "bsonType": "array", "items": { "bsonType": "object", "properties": { "provider": {"enum": ["jazzcash", "easypaisa", "raast", "nayapay", "meezan"]}, "account_id": {"bsonType": "string"} } } }, "is_active": {"bsonType": "bool"} } } } try: await db.create_collection("users", validator=user_validator) except Exception: pass await db.users.create_indexes([ IndexModel([("phone", ASCENDING)], unique=True), IndexModel([("user_id", ASCENDING)], unique=True) ]) # --- 2. SIGNALS COLLECTION --- logger.info("Setting up 'signals' collection...") try: await db.create_collection("signals") except Exception: pass await db.signals.create_indexes([ IndexModel([("user_id", ASCENDING)]), IndexModel([("submitted_at", DESCENDING)]) ]) # --- 3. VERDICTS COLLECTION --- logger.info("Setting up 'verdicts' collection...") try: await db.create_collection("verdicts") except Exception: pass await db.verdicts.create_indexes([ IndexModel([("signal_id", ASCENDING)]), IndexModel([("user_id", ASCENDING)]), IndexModel([("risk_level", ASCENDING)]) ]) # --- 4. FRAUD LEDGER COLLECTION --- logger.info("Setting up 'fraud_ledger' collection...") try: await db.create_collection("fraud_ledger") except Exception: pass await db.fraud_ledger.create_indexes([ IndexModel([("sender_hash", ASCENDING)], unique=True), IndexModel([("url_hash", ASCENDING)]), IndexModel([("scam_pattern", TEXT)]), IndexModel([("hit_count", DESCENDING)]) ]) # --- 5. COMPLAINTS COLLECTION --- logger.info("Setting up 'complaints' collection...") try: await db.create_collection("complaints") except Exception: pass await db.complaints.create_indexes([ IndexModel([("user_id", ASCENDING)]), IndexModel([("fia_ref", ASCENDING)], unique=True), IndexModel([("status", ASCENDING)]) ]) # --- 6. WALLET HOLDS COLLECTION --- logger.info("Setting up 'wallet_holds' collection...") try: await db.create_collection("wallet_holds") except Exception: pass await db.wallet_holds.create_indexes([ IndexModel([("user_id", ASCENDING)]), IndexModel([("is_active", ASCENDING)]), IndexModel([("expires_at", ASCENDING)]) ]) # --- 7. TRACES COLLECTION --- logger.info("Setting up 'traces' collection...") try: await db.create_collection("traces") except Exception: pass await db.traces.create_indexes([ IndexModel([("workplan_id", ASCENDING)], unique=True), IndexModel([("signal_id", ASCENDING)]), IndexModel([("created_at", DESCENDING)]) ]) # --- 8. ALERTS COLLECTION --- logger.info("Setting up 'alerts' collection...") try: await db.create_collection("alerts") except Exception: pass await db.alerts.create_indexes([ IndexModel([("signal_id", ASCENDING)]), IndexModel([("sent_by_user_id", ASCENDING)]) ]) # --- SEED DATA --- logger.info("Inserting seed data for demo...") # Clean old seed data for a fresh setup await db.users.delete_many({}) await db.fraud_ledger.delete_many({}) # Seed Users users_seed = [ { "user_id": str(uuid.uuid4()), "phone": "0312-3456789", "language_preference": "roman_ur", "family_circle": [{"name": "Beta Kamran", "phone": "0333-1234567", "language": "ur", "relationship": "son"}], "wallet_accounts": [{"provider": "meezan", "account_id": "PK77MEZN0001"}], "is_active": True, "created_at": datetime.now() }, { "user_id": str(uuid.uuid4()), "phone": "0321-9876543", "language_preference": "ur", "family_circle": [], "wallet_accounts": [{"provider": "jazzcash", "account_id": "03219876543"}], "is_active": True, "created_at": datetime.now() } ] await db.users.insert_many(users_seed) # Seed Fraud Ledger (Pakistani Scam Patterns) ledger_seed = [ {"sender_hash": "hash_jazz_help", "scam_pattern": "JazzCash Fake Helpline Calling from 0300 numbers", "scam_type": "fake_helpline", "hit_count": 47, "reported_by_count": 12}, {"sender_hash": "hash_ep_otp", "scam_pattern": "EasyPaisa OTP request for 'Wrong Transfer' refund", "scam_type": "smishing", "hit_count": 23, "reported_by_count": 9}, {"sender_hash": "hash_olx", "scam_pattern": "OLX Advance Payment for Bike or iPhone", "scam_type": "olx_fraud", "hit_count": 31, "reported_by_count": 15}, {"sender_hash": "hash_wa_imp", "scam_pattern": "WhatsApp Impersonation 'Assalam o Alaikum, Main mushkil mein hoon'", "scam_type": "whatsapp_impersonation", "hit_count": 15, "reported_by_count": 6}, {"sender_hash": "hash_loan", "scam_pattern": "Urgent Cash Loan App Blackmail/Extortion", "scam_type": "loan_extortion", "hit_count": 8, "reported_by_count": 4}, {"sender_hash": "hash_benazir", "scam_pattern": "BISP / Benazir Income Support Program fake SMS", "scam_type": "smishing", "hit_count": 55, "reported_by_count": 20}, {"sender_hash": "hash_jeeto", "scam_pattern": "Jeeto Pakistan / Inaami Scheme fraud message", "scam_type": "investment_scam", "hit_count": 12, "reported_by_count": 3}, {"sender_hash": "hash_hbl", "scam_pattern": "HBL Account Blocked fake link phishing", "scam_type": "smishing", "hit_count": 19, "reported_by_count": 7}, {"sender_hash": "hash_lottery", "scam_pattern": "Saudi Arabia Lottery fake WhatsApp call", "scam_type": "investment_scam", "hit_count": 5, "reported_by_count": 2}, {"sender_hash": "hash_fbr", "scam_pattern": "FBR Tax Refund fake website scam", "scam_type": "unknown", "hit_count": 9, "reported_by_count": 4} ] await db.fraud_ledger.insert_many(ledger_seed) await verify_setup(db) logger.info("NadraGuard DB setup complete.") async def verify_setup(db): logger.info("--- VERIFICATION SUMMARY ---") collections = await db.list_collection_names() for col in collections: count = await db[col].count_documents({}) indexes = await db[col].list_indexes().to_list(length=100) logger.info(f"Collection: {col.ljust(15)} | Documents: {str(count).ljust(5)} | Indexes: {len(indexes)}") if __name__ == "__main__": asyncio.run(setup_database())