Merge pull request #19 from Gurummang/develop

feat: add all-scan-queue

Author: buubb

app/__init__.py

@@ -40,12 +40,14 @@
 DOC_SCAN_QUEUE = os.getenv("RABBITMQ_DOC_QUEUE_NAME")
 EXE_SCAN_QUEUE = os.getenv("RABBITMQ_EXE_QUEUE_NAME")
 IMG_SCAN_QUEUE = os.getenv("RABBITMQ_IMG_QUEUE_NAME")
+ALL_SCAN_QUEUE = os.getenv("RABBITMQ_ALL_QUEUE_NAME")
 ALERT_QUEUE = os.getenv("RABBITMQ_SUSPICIOUS_QUEUE")
 
 # Routing Key 설정
 EXE_ROUTING_KEY = os.getenv("RABBITMQ_EXE_ROUTING_KEY")
 IMG_ROUTING_KEY = os.getenv("RABBITMQ_IMG_ROUTING_KEY")
 DOC_ROUTING_KEY = os.getenv("RABBITMQ_DOC_ROUTING_KEY")
+ALL_ROUTING_KEY = os.getenv("RABBITMQ_ALL_ROUTING_KEY")
 ALERT_ROUTING_KEY = os.getenv("RABBITMQ_SUSPICIOUS_ROUTING_KEY")
 
 # S3

app/utils.py

@@ -21,6 +21,31 @@
 from app.models import FileScanRequest
 from app.rabbitmq_sender import send_message
 
+def match_multiple_rules(*rule_file_lists):
+    """
+    여러 YARA 룰셋의 파일 경로를 병합하고 컴파일된 하나의 룰셋을 반환하는 함수.
+    """
+    all_rule_files = {}
+    rule_index = 0
+    
+    # 각 룰셋의 파일 경로 리스트에서 유효한 룰 파일들을 수집
+    for rule_files in rule_file_lists:
+        for i, rule_file in enumerate(rule_files):
+            all_rule_files[str(rule_index + i)] = rule_file
+        rule_index += len(rule_files)
+
+    if all_rule_files:
+        try:
+            # 여러 YARA 룰 파일을 하나의 룰셋으로 컴파일
+            compiled_rules = yara.compile(filepaths=all_rule_files)
+            return compiled_rules
+        except yara.Error as e:
+            logging.error(f"Failed to compile merged YARA rules: {e}")
+            return None
+    else:
+        logging.info("No valid YARA rule files found for merging.")
+        return None
+
 
 def load_yara_rules(directory):
     rule_files = []
@@ -30,7 +55,6 @@ def load_yara_rules(directory):
         logging.info(f"Scanning directory: {root}")  # 현재 디렉토리 로그에 남기기
         for file in files:
             if file.endswith(".yar"):
-                logging.info(f"Scanning file: {file}")  # 현재 디렉토리 로그에 남기기
                 rule_files.append(os.path.join(root, file))
 
     # YARA 룰 컴파일
@@ -54,16 +78,17 @@ def load_yara_rules(directory):
                 logging.info(
                     f"Compiled {len(valid_rule_files)} YARA rules from {directory}"
                 )
-                return compiled_rules
+                return compiled_rules, valid_rule_files
             except yara.Error as e:
                 logging.info(f"Failed to compile YARA rules: {e}")
-                return None
+                return None, valid_rule_files
         else:
             logging.info(f"No valid YARA rule files found in {directory}")
-            return None
+            return None, []
     else:
         logging.info(f"No YARA rule files found in {directory}")
-        return None
+        return None, []
+
 
 
 def stream_file_from_s3(s3_key):

main.py

@@ -14,9 +14,11 @@
     EXE_SCAN_QUEUE,
     IMG_ROUTING_KEY,
     IMG_SCAN_QUEUE,
+    ALL_SCAN_QUEUE,
+    ALL_ROUTING_KEY
 )
 from app.rabbitmq_consumer import start_consuming
-from app.utils import load_yara_rules
+from app.utils import load_yara_rules, match_multiple_rules
 
 app = FastAPI()
 
@@ -33,9 +35,9 @@ async def lifespan(app: FastAPI):
 
     try:
         # YARA 규칙을 로드하고 컴파일
-        rules["exe"] = load_yara_rules(os.path.join(RULES_DIR, "exe"))
-        rules["img"] = load_yara_rules(os.path.join(RULES_DIR, "img"))
-        rules["doc"] = load_yara_rules(os.path.join(RULES_DIR, "doc"))
+        rules["exe"], exe_files = load_yara_rules(os.path.join(RULES_DIR, "exe"))
+        rules["img"], img_files = load_yara_rules(os.path.join(RULES_DIR, "img"))
+        rules["doc"], doc_files = load_yara_rules(os.path.join(RULES_DIR, "doc"))
         logging.info("YARA rules loaded and compiled successfully.")
     except Exception as e:
         logging.error(f"Failed to load YARA rules: {e}")
@@ -49,6 +51,12 @@ async def lifespan(app: FastAPI):
     Thread(
         target=start_consuming, args=(DOC_SCAN_QUEUE, rules["doc"], DOC_ROUTING_KEY)
     ).start()
+    # ALL QUEUE에는 모든 규칙 적용
+    all_rules_matcher = match_multiple_rules(doc_files, exe_files, img_files)
+    Thread(
+        target=start_consuming, 
+        args=(ALL_SCAN_QUEUE, all_rules_matcher, ALL_ROUTING_KEY)
+    ).start()
 
     yield