Merge pull request #9 from Gurummang/develop

KAN-456 <feat: set RabbitMQ for alerts>

Author: buubb

app/__init__.py

@@ -34,16 +34,19 @@
 # Exchange 설정
 EXCHANGE_NAME = os.getenv("RABBITMQ_EXCHANGE_NAME")
 EXCHANGE_TYPE = os.getenv("RABBITMQ_EXCHANGE_TYPE")
+ALERT_EXCHANGE_NAME = os.getenv("RABBITMQ_ALERT_EXCHANGE")
 
 # Queue 설정
 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")
+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")
+ALERT_ROUTING_KEY = os.getenv("RABBITMQ_SUSPICIOUS_ROUTING_KEY")
 
 # S3
 S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")

app/models.py

@@ -1,5 +1,5 @@
 from pydantic import BaseModel
-from sqlalchemy import BigInteger, Boolean, Column, Integer, Text
+from sqlalchemy import BigInteger, Boolean, Column, Integer, Text, String, DateTime
 from sqlalchemy.ext.declarative import declarative_base
 
 Base = declarative_base()
@@ -34,3 +34,13 @@ class FileStatus(Base):
     gscan_status = Column(Boolean, nullable=True, default=-1)
     dlp_status = Column(Boolean, nullable=True, default=-1)
     vt_status = Column(Boolean, nullable=True, default=-1)
+
+
+class FileUpload(Base):
+    __tablename__ = "file_upload"
+    id = Column(BigInteger, primary_key=True, autoincrement=True)
+    org_saas_id = Column(Integer, nullable=False)
+    saas_file_id = Column(String(255), nullable=True)
+    upload_ts = Column(DateTime, nullable=True)
+    salted_hash = Column(String(255), nullable=True)
+    deleted = Column(Boolean, default=False, nullable=False)

app/rabbitmq_sender.py

@@ -0,0 +1,70 @@
+import logging
+import ssl
+import time
+import struct
+import pika
+
+from app import (
+    ALERT_EXCHANGE_NAME,
+    ALERT_ROUTING_KEY,
+    EXCHANGE_TYPE,
+    RABBITMQ_HOST,
+    RABBITMQ_PASSWORD,
+    RABBITMQ_PORT,
+    RABBITMQ_SSL_ENABLED,
+    RABBITMQ_USER,
+    RETRY_INTERVAL,
+)
+
+# SSL 설정
+ssl_options = None
+if RABBITMQ_SSL_ENABLED:
+    ssl_context = ssl.create_default_context(purpose=ssl.Purpose.SERVER_AUTH)
+    ssl_options = pika.SSLOptions(context=ssl_context)
+
+
+def connect_to_rabbitmq():
+    while True:
+        try:
+            credentials = pika.PlainCredentials(RABBITMQ_USER, RABBITMQ_PASSWORD)
+            parameters = pika.ConnectionParameters(
+                host=RABBITMQ_HOST,
+                port=int(RABBITMQ_PORT),
+                credentials=credentials,
+                ssl_options=ssl_options,
+                connection_attempts=3,
+                retry_delay=5,
+                socket_timeout=10.0,  # 타임아웃 설정 (초)
+            )
+            connection = pika.BlockingConnection(parameters)
+            return connection
+        except pika.exceptions.AMQPConnectionError as e:
+            logging.error(
+                f"Connection failed, retrying in {RETRY_INTERVAL} seconds... Error: {e}"
+            )
+            time.sleep(RETRY_INTERVAL)
+
+
+def send_message(message: int):
+    connection = connect_to_rabbitmq()
+    channel = connection.channel()
+
+    # Exchange 선언
+    channel.exchange_declare(
+        exchange=ALERT_EXCHANGE_NAME, exchange_type=EXCHANGE_TYPE, durable=True
+    )
+
+    # int 메시지를 바이트로 변환
+    message_bytes = struct.pack('!Q', message)  # '!Q'는 unsigned long long 형식입니다.
+
+    channel.basic_publish(
+        exchange=ALERT_EXCHANGE_NAME,
+        routing_key=ALERT_ROUTING_KEY,  # 적절한 라우팅 키로 변경
+        body=message_bytes,
+        properties=pika.BasicProperties(
+            delivery_mode=2  # 메시지 영속화
+        )
+    )
+
+    print(f"Sent message: {message}")
+    connection.close()

app/utils.py

@@ -16,7 +16,7 @@
     MYSQL_USER,
 )
 from app.models import FileScanRequest
-
+from app.rabbitmq_sender import send_message
 
 def load_yara_rules(directory):
     rule_files = []
@@ -74,7 +74,7 @@ def stream_file_from_s3(s3_key):
         raise
 
 
-def save_scan_result(file_id, detect, detail):
+def save_scan_result(uploadId: int, stored_file_id, detect, detail):
     try:
         conn = mysql.connector.connect(
             host=MYSQL_HOST, user=MYSQL_USER, password=MYSQL_PASSWORD, database=MYSQL_DB
@@ -84,7 +84,7 @@ def save_scan_result(file_id, detect, detail):
         try:
             cursor.execute(
                 "INSERT INTO scan_table (file_id, detect, step2_detail) VALUES (%s, %s, %s)",
-                (file_id, detect, detail),
+                (stored_file_id, detect, detail),
             )
             conn.commit()  # 첫 번째 쿼리 커밋
         except Exception as e:
@@ -94,9 +94,10 @@ def save_scan_result(file_id, detect, detail):
 
         try:
             cursor.execute(
-                "UPDATE file_status SET gscan_status = 1 WHERE file_id = %s", (file_id,)
+                "UPDATE file_status SET gscan_status = 1 WHERE file_id = %s", (stored_file_id,)
             )
             conn.commit()  # 두 번째 쿼리 커밋
+            send_message(uploadId) # RabbitMQ 전송: Alerts
         except Exception as e:
             conn.rollback()  # 두 번째 쿼리 롤백
             logging.error(f"Failed to update file_status: {e}")
@@ -147,21 +148,17 @@ def yara_test_match(file_path, yara_rules):
     return detect, detail
 
 
-def scan_file(file_id: int, yara_rules):
+def scan_file(upload_id: int, yara_rules):
     try:
-        conn = mysql.connector.connect(
-            host=MYSQL_HOST, user=MYSQL_USER, password=MYSQL_PASSWORD, database=MYSQL_DB
-        )
-        cursor = conn.cursor(dictionary=True)
-        cursor.execute("SELECT * FROM stored_file WHERE id = %s", (file_id,))
-        file_record = cursor.fetchone()
-        cursor.close()
-        conn.close()
-
-        if not file_record:
-            raise HTTPException(status_code=404, detail="File not found")
-
-        s3_key = file_record["save_path"]
+        # 파일 업로드 정보 가져오기
+        file_record = get_file_upload(upload_id)
+        salted_hash = file_record["salted_hash"]
+
+        # 저장된 파일 정보 가져오기
+        stored_file_record = get_stored_file(salted_hash)
+        stored_file_id = stored_file_record["id"]
+        s3_key = stored_file_record["save_path"]
+    
         file_stream = stream_file_from_s3(s3_key)
 
         # 파일 전체를 한 번에 읽음
@@ -182,7 +179,49 @@ def scan_file(file_id: int, yara_rules):
         logging.info(f"detail: {detail}")
         logging.info(f"most_common_keyword: {most_common_keyword}")
 
-        save_scan_result(file_id, detect, most_common_keyword)
+        save_scan_result(upload_id, stored_file_id, detect, most_common_keyword)
     except Exception as e:
         logging.error(f"Error scanning file: {e}")
         raise HTTPException(status_code=500, detail="Error scanning file")
+    
+
+def get_stored_file(hash:str):
+    try:
+        conn = mysql.connector.connect(
+            host=MYSQL_HOST, user=MYSQL_USER, password=MYSQL_PASSWORD, database=MYSQL_DB
+        )
+        cursor = conn.cursor(dictionary=True)
+        cursor.execute("SELECT * FROM stored_file WHERE salted_hash = %s", (hash,))
+        stored_file_record = cursor.fetchone()
+        cursor.close()
+        conn.close()
+
+        if not stored_file_record:
+            raise HTTPException(status_code=404, detail="File not found in stored_file table")
+
+        return stored_file_record
+
+    except Exception as e:
+        logging.error(f"Error fetching stored file record: {e}")
+        raise HTTPException(status_code=500, detail="Error fetching stored file record")
+
+
+def get_file_upload(file_id: int):
+    try:
+        conn = mysql.connector.connect(
+            host=MYSQL_HOST, user=MYSQL_USER, password=MYSQL_PASSWORD, database=MYSQL_DB
+        )
+        cursor = conn.cursor(dictionary=True)
+        cursor.execute("SELECT * FROM file_upload WHERE id = %s", (file_id,))
+        file_record = cursor.fetchone()
+        cursor.close()
+        conn.close()
+
+        if not file_record:
+            raise HTTPException(status_code=404, detail="File not found in file_upload table")
+
+        return file_record
+
+    except Exception as e:
+        logging.error(f"Error fetching file upload record: {e}")
+        raise HTTPException(status_code=500, detail="Error fetching file upload record")