dlp_proxy.py

import json
import os
import re
from mitmproxy import http
from mitmproxy import ctx
from presidio_analyzer import AnalyzerEngine
from typing import List
from modules.check_regex import checkRegex
from modules.check_keyword import checkKeyword
from modules.check_upload import checkUploadGithub
from modules.handle_sensitive_data import handleSensitiveData

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

CONFIG_DIR = os.path.join(BASE_DIR, "json")
RULE_FILE = os.path.join(CONFIG_DIR, "sensitive_info.json")
KEYWORD_FILE = os.path.join(CONFIG_DIR, "keyword.json")
ALLOWLIST_FILE = os.path.join(CONFIG_DIR, "allowlist.json")
LOG_FILE = "dlp.log"

# 허용되는 최대 크기 100MB
MAX_PAYLOAD_SIZE = 100 * 1024 * 1024

class EnglishDLPAnalyzer:
    def __init__(self):
        self.analyzer = AnalyzerEngine()

    def analyze_text(self, text: str):
        """
        Returns:
            {
                "is_leak": bool,
                "entities": ["EMAIL_ADDRESS", "PHONE_NUMBER", ...]
            }
        """
        if not text or text.strip() == "":
            return {"is_leak": False, "entities": []}

        results = self.analyzer.analyze(text=text, language="en")

        if not results:
            return {"is_leak": False, "entities": []}

        entities = list(set([r.entity_type for r in results]))

        return {
            "is_leak": True,
            "entities": entities
        }

nlp_dlp = EnglishDLPAnalyzer()

try:
    with open(RULE_FILE, "r", encoding="utf-8") as f:
        DLP_RULES = json.load(f)
        ctx.log.info(f"[*] Loaded {len(DLP_RULES)} DLP rules from {RULE_FILE}")
except Exception as e:
    ctx.log.error(f"[!] Failed to load rules: {e}")
    DLP_RULES = []

try:
    with open(KEYWORD_FILE, "r", encoding="utf-8") as f:
        KEYWORD = json.load(f)
        ctx.log.info(f"[*] Loaded {len(KEYWORD)} keyword from {KEYWORD_FILE}")
except Exception as e:
    ctx.log.error(f"[!] Failed to load keyword: {e}")
    KEYWORD = []

try:
    with open(ALLOWLIST_FILE, "r", encoding="utf-8") as f:
        ALLOWLIST = json.load(f)
        ctx.log.info(f"[*] Loaded {len(ALLOWLIST)} keyword from {ALLOWLIST_FILE}")
except Exception as e:
    ctx.log.error(f"[!] Failed to load keyword: {e}")
    ALLOWLIST = []

def request(flow: http.HTTPFlow) -> None:
    request_url = flow.request.pretty_url
    request_header = json.dumps(dict(flow.request.headers))
    request_body = flow.request.get_text()
    request_method = flow.request.method
    request_text = f"URL : {request_url}, Headers : {request_header}, Body : {request_body}"
    content_type = flow.request.headers.get("Content-Type", "")

    ALLOWLIST_DOMAINS = ALLOWLIST.get("ALLOWLIST_DOMAINS")

    if ALLOWLIST_DOMAINS and re.search(ALLOWLIST_DOMAINS, request_url):
        return

    if any(x in content_type for x in ["video", "audio"]):
        return

    if not request_text:
        return

    detected_items = checkRegex(request_text, DLP_RULES, KEYWORD)
    if detected_items:
        handleSensitiveData(flow, detected_items, request_url)
    
    detected_items = checkKeyword(request_url, request_text, KEYWORD, ALLOWLIST, MAX_PAYLOAD_SIZE)
    if detected_items:
        handleSensitiveData(flow, detected_items, request_url)

    detected_items = checkUploadGithub(request_url, request_method, request_header, ALLOWLIST)
    if detected_items:
        handleSensitiveData(flow, detected_items, request_url)
    
    ai_result = nlp_dlp.analyze_text(request_body)

    if ai_result["is_leak"]:
        alert_msg = f"[DLP DETECTED/AI] {request_url} -> Entities: {ai_result['entities']}"
        ctx.log.alert(alert_msg)

        with open(LOG_FILE, "a", encoding="utf-8") as log:
            log.write(f"{alert_msg}\n")
            log.write(f"Message request payload: {request_text}\n")

        flow.response = http.Response.make(
            403,
            b"Blocked by AI-DLP: Sensitive Information Found\n",
            {"Content-Type": "text/plain"}
        )
        return
    
def response(flow: http.HTTPFlow) -> None:
    response_header = json.dumps(dict(flow.response.headers))
    response_body = flow.response.get_text()
    response_text = f"Headers : {response_header}, Body : {response_body}"
    content_type = flow.response.headers.get("Content-Type", "")
    if any(x in content_type for x in ["video", "audio"]):
        return
    
    if not response_text:
        return

    detected_items = checkRegex(response_text, DLP_RULES, KEYWORD)
    if detected_items:
        handleSensitiveData(flow, detected_items)
    
    
    ai_result = nlp_dlp.analyze_text(response_text)

    if ai_result["is_leak"]:
        alert_msg = f"[DLP DETECTED/AI] {flow.request.pretty_url} -> Entities: {ai_result['entities']}"
        ctx.log.alert(alert_msg)

        with open(LOG_FILE, "a", encoding="utf-8") as log:
            log.write(f"{alert_msg}\n")
            log.write(f"Message response payload: {response_text}\n")

        flow.response = http.Response.make(
            403,
            b"Blocked by AI-DLP: Sensitive Information Found\n",
            {"Content-Type": "text/plain"}
        )
        return