Add yara rules

app.py

@@ -9,9 +9,12 @@ import shutil
 import cv2
 import logging
 import uuid
+import yara
+import asyncio
 from huggingface_hub import snapshot_download
-from typing import Optional
+from typing import Optional, List, Dict
 import aiofiles
+from fastapi.concurrency import run_in_threadpool
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -25,25 +28,103 @@ MAL_CLASSES = ['Adialer.C', 'Agent.FYI', 'Allaple.A', 'Allaple.L', 'Alueron.gen!
 
 UPLOAD_DIR = "uploads"
 os.makedirs(UPLOAD_DIR, exist_ok=True)
+YARA_REPO_URL = "https://github.com/roadwy/DefenderYara.git"
+YARA_REPO_DIR = "defenderyara"
 
 # Environment configuration
 MODEL_REPO = os.getenv("MODEL_REPO", "GranularFireplace/malware")
 MODEL_FILE = os.getenv("MODEL_FILE", "model_v2_with_weight.keras")
 
+async def clone_yara_repo():
+    """Clone YARA rules repository asynchronously"""
+    try:
+        repo_path = Path(YARA_REPO_DIR)
+        
+        # Remove existing repository if it exists
+        if repo_path.exists():
+            logger.info("Removing existing YARA rules repository")
+            shutil.rmtree(repo_path)
+
+        logger.info(f"Cloning YARA rules from {YARA_REPO_URL}")
+        proc = await asyncio.create_subprocess_exec(
+            'git', 'clone', YARA_REPO_URL, str(repo_path),
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE
+        )
+
+        stdout, stderr = await proc.communicate()
+
+        if proc.returncode != 0:
+            logger.error(f"Failed to clone YARA repo: {stderr.decode()}")
+            return None
+        
+        logger.info("YARA rules repository cloned successfully")
+        return repo_path
+
+    except Exception as e:
+        logger.error(f"Error cloning YARA repository: {str(e)}")
+        return None
+
+def compile_yara_rules(repo_path: Path) -> Optional[yara.Rules]:
+    """Compile YARA rules from repository"""
+    try:
+        yara_files = list(repo_path.glob("**/*.yar"))
+        if not yara_files:
+            logger.warning("No YARA files found in repository")
+            return None
+
+        logger.info(f"Found {len(yara_files)} YARA files, compiling rules")
+        rules = {}
+        
+        for yara_file in yara_files:
+            try:
+                rules[str(yara_file)] = str(yara_file)
+            except Exception as e:
+                logger.warning(f"Error processing {yara_file}: {str(e)}")
+
+        return yara.compile(filepaths=rules)
+    
+    except yara.SyntaxError as e:
+        logger.error(f"YARA syntax error: {str(e)}")
+    except Exception as e:
+        logger.error(f"Error compiling YARA rules: {str(e)}")
+    
+    return None
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """Manage model loading and unloading during app lifecycle"""
+    """Manage application lifecycle"""
+    # Initialize app state
+    app.state.model = None
+    app.state.yara_rules = None
+
     try:
+        # Load ML model
         logger.info("Downloading model from Hugging Face Hub...")
         download_dir = snapshot_download(MODEL_REPO)
         app.state.model = tf.keras.models.load_model(os.path.join(download_dir, MODEL_FILE))
         logger.info("Model loaded successfully")
+
+        # Clone and compile YARA rules
+        yara_repo_path = await clone_yara_repo()
+        if yara_repo_path:
+            app.state.yara_rules = compile_yara_rules(yara_repo_path)
+            if app.state.yara_rules:
+                logger.info("YARA rules compiled successfully")
+            else:
+                logger.warning("No valid YARA rules compiled")
+        else:
+            logger.warning("YARA rules unavailable")
+
     except Exception as e:
-        logger.error(f"Error loading model: {str(e)}")
+        logger.error(f"Initialization error: {str(e)}")
         raise
+
     yield
-    # Cleanup resources if needed
+
+    # Cleanup
     app.state.model = None
+    app.state.yara_rules = None
 
 app = FastAPI(lifespan=lifespan)