proxy_script.py

import sys
from mitmproxy import http
from transformers import pipeline
import logging
from translate import Translator
from database import add_blocked_request

# Для Windows: отключаем использование urwid
if sys.platform == "win32":
    from mitmproxy.tools.main import mitmdump

    sys.argv.append("--no-web")

# Инициализация модели
model_name = "./rubert-sentiment"
classifier = pipeline(
    "sentiment-analysis",
    model=model_name,
    device=-1  # -1 - CPU, 0 - GPU
)

logging.basicConfig(
    filename='proxy.log',
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)


def is_malicious(text):
    # translator = Translator(from_lang='russian', to_lang='english')
    # translation = translator.translate(text)
    if not text.strip():
        return False

    result = classifier(text[:512], truncation=True)[0]
    return result['label'] == 'NEGATIVE' and result['score'] > 0.7


def request(flow: http.HTTPFlow) -> None:
    if flow.request.method in ('GET', 'POST'):
        url = flow.request.pretty_url
        content = ""

        # Извлечение контента
        if flow.request.method == 'GET':
            content = flow.request.query.get('q', '') or flow.request.query.get('query', '')
        elif flow.request.method == 'POST':
            if 'application/x-www-form-urlencoded' in flow.request.headers.get('Content-Type', ''):
                content = flow.request.urlencoded_form.get('text', '') or flow.request.urlencoded_form.get('search', '')
            elif 'application/json' in flow.request.headers.get('Content-Type', ''):
                try:
                    content = flow.request.json().get('query', '') or flow.request.json().get('text', '')
                except:
                    pass

        if content and is_malicious(content):
            # add_blocked_request(
            #     url,
            #     flow.request.method,
            #     str(flow.request.headers),
            #     content,
            #     "Негативный контент"
            # )
            flow.response = http.Response.make(
                403,
                "Запрос заблокирован".encode('utf-8'),
                {"Content-Type": "text/plain; charset=utf-8"}
            )
            # logging.warning(f"Блокировка: {url} | Контент: {content}")