↓Перейти к основному содержимому

ИИ Антиспам на Proxmox Mail Gateway

·2820 слов·14 минут
Андрей Двуреченский
Автор
Андрей Двуреченский

Введение
#

Proxmox Mail Gateway - шлюз для электронной почты, который стоит до почтового сервера, фильтрует спам, вирусы и пр.

Про настройку самого Proxmox Mail Gateway документации и видео в интернете достаточно (например чать 1, чать 2).

Custom check
#

Почему бы дополнительно к спам-фильтрам не добавить AI, подумал я и решил этим заняться. Напрямую интеграции PMG с какими-либо нейросетями нет, но в документации есть упоминание функции custom check - это скрипт, который выполняется еще до фильтров самого PMG.

Вайб-кодим небольшой скрипт, который будет наши письма анализировать с помощью AI. В качестве провайдера - я использую OpenRouter и небольшую модель mistralai/mistral-small-3.1-24b-instruct.

Используем white-лист чтобы не анализировать доверенные домены и black-лист чтобы повторно не проверять письма с домена, который уже попался на рассылке спама.

#!/bin/bash

# Custom email checker script for Proxmox Mail Gateway
# Uses OpenRouter API to analyze emails for spam content

set -euo pipefail  # Exit on error, undefined variables, and pipe failures

# Log file configuration
LOG_FILE="/var/log/pmg-custom-check.log"
LOG_MAX_SIZE=10485760  # 10MB

# Rotate log file if it exceeds max size
rotate_log() {
    if [ -f "$LOG_FILE" ]; then
        local size=$(stat -c%s "$LOG_FILE" 2>/dev/null || echo 0)
        if [ "$size" -gt "$LOG_MAX_SIZE" ]; then
            mv "$LOG_FILE" "${LOG_FILE}.old" 2>/dev/null || true
        fi
    fi
}

# Log function to write to both STDERR (mail.log) and log file
log() {
    local message
    message="$(date '+%Y-%m-%d %H:%M:%S') - $1"
    # Explicitly write to file descriptor 2 (stderr) and handle any errors
    {
        printf '%s\n' "$message" >&2
    } 2>/dev/null || true
    rotate_log
    {
        printf '%s\n' "$message" >> "$LOG_FILE"
    } 2>/dev/null || true
}

# Function to extract email headers
extract_email_headers() {
    local email_file="$1"
    
    # Extract sender (From header) and decode base64 encoding if present
    local raw_sender=$(grep -m 1 '^From:' "$email_file" 2>/dev/null | sed 's/^From: *//' || echo "unknown")
    
    # Decode MIME encoded words (e.g., =?UTF-8?B?...?= or =?UTF-8?Q?...?=)
    local sender="$raw_sender"
    if [[ "$raw_sender" == *=?* ]]; then
        # Use Python for more reliable MIME decoding
        if command -v python3 >/dev/null 2>&1; then
            # Create temporary file approach to avoid complex quoting
            local temp_py_script="/tmp/pmg_decode_$$.py"
            cat > "$temp_py_script" << 'EOF_PY'
import sys
from email.header import decode_header
raw_sender = sys.argv[1] if len(sys.argv) > 1 else ""
decoded = decode_header(raw_sender)
result = []
for part, encoding in decoded:
    if isinstance(part, bytes):
        if encoding:
            result.append(part.decode(encoding))
        else:
            result.append(part.decode('utf-8'))
    else:
        result.append(part)
print(''.join(result))
EOF_PY
            # Ensure the Python script executes safely
            if [ -f "$temp_py_script" ]; then
                sender=$(python3 "$temp_py_script" "$raw_sender" 2>/dev/null)
                # If Python decoding failed, use original value
                if [ $? -ne 0 ] || [ -z "$sender" ]; then
                    sender="$raw_sender"
                fi
                rm -f "$temp_py_script"
            else
                # Fallback if temp script creation failed
                sender="$raw_sender"
            fi
        fi
    fi
    
    # Extract receiver (To header)
    local receiver=$(grep -m 1 '^To:' "$email_file" 2>/dev/null | sed 's/^To: *//' || echo "unknown")
    
    # Extract subject (Subject header)
    local subject=$(grep -m 1 '^Subject:' "$email_file" 2>/dev/null | sed 's/^Subject: *//' || echo "no subject")
    
    # Extract message ID (Message-ID header) - keeping for internal use only
    local message_id=$(grep -m 1 '^Message-ID:' "$email_file" 2>/dev/null | sed 's/^Message-ID: *//' || echo "no message-id")
    
    # Print the extracted headers for return
    echo "$sender|$receiver|$subject|$message_id"
}

# Print usage and exit if arguments are incorrect
if [ "$#" -ne 2 ]; then
    log "usage: $0 APIVERSION QUEUEFILENAME"
    exit 1
fi

apiver="$1"
queue_file="$2"

# Check API version
if [ "$apiver" != "v1" ]; then
    log "wrong APIVERSION: $apiver"
    exit 2
fi

# Extract email headers before processing (needed for logging in error cases)
headers_info=$(extract_email_headers "$queue_file")
sender=$(echo "$headers_info" | cut -d'|' -f1)
receiver=$(echo "$headers_info" | cut -d'|' -f2)
subject=$(echo "$headers_info" | cut -d'|' -f3)
message_id=$(echo "$headers_info" | cut -d'|' -f4)

# Whitelist/Blacklist functionality
check_whitelist_blacklist() {
    local sender="$1"
    local whitelist_file="$(dirname "$0")/whitelist.txt"
    local blacklist_file="$(dirname "$0")/blacklist.txt"
    
    # Check if sender matches whitelist
    if [ -f "$whitelist_file" ]; then
        while IFS= read -r pattern; do
            # Skip empty lines
            [ -z "$pattern" ] && continue
            if echo "$sender" | grep -qi "$pattern"; then
                log "Sender in whitelist: $sender - returning SCORE: 0"
                {
                    printf '%s\n' "SCORE: 0" >&1
                } 2>/dev/null || true
                exit 0
            fi
        done < "$whitelist_file"
    fi
    
    # Check if sender matches blacklist
    if [ -f "$blacklist_file" ]; then
        while IFS= read -r pattern; do
            # Skip empty lines
            [ -z "$pattern" ] && continue
            if echo "$sender" | grep -qi "$pattern"; then
                log "Sender in blacklist: $sender - returning SCORE: 10"
                {
                    printf '%s\n' "SCORE: 10" >&1
                } 2>/dev/null || true
                exit 0
            fi
        done < "$blacklist_file"
    fi
}

# Check whitelist/blacklist before proceeding with API check
check_whitelist_blacklist "$sender"

# Check if input file exists
if [ ! -f "$queue_file" ]; then
    log "queue file does not exist: $queue_file"
    # Output API version and result for missing file to STDOUT explicitly
    {
        printf '%s\n' "v1" >&1
        printf '%s\n' "OK" >&1
    } 2>/dev/null || true
    exit 0
fi

# Configuration for OpenRouter API
OPENROUTER_API_KEY="sk-or-v1-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
OPENROUTER_MODEL="mistralai/mistral-small-3.1-24b-instruct"

if [ -z "$OPENROUTER_API_KEY" ]; then
    log "ERROR: OPENROUTER_API_KEY environment variable not set, cannot proceed with API check - Sender: $sender, Receiver: $receiver"
    # Output API version and result if no API key to STDOUT explicitly
    {
        printf '%s\n' "v1" >&1
        printf '%s\n' "OK" >&1
    } 2>/dev/null || true
    exit 0
fi

# Output API version once all validation is done to STDOUT explicitly
{
    printf '%s\n' "v1" >&1
} 2>/dev/null || true

# Read the email content
email_content=""
if [ -r "$queue_file" ]; then
    email_content=$(cat "$queue_file")
    log "Successfully read email content - Sender: $sender, Receiver: $receiver, Size: $(echo -n "$email_content" | wc -c) bytes"
else
    log "ERROR: cannot read queue file: $queue_file - Sender: $sender, Receiver: $receiver"
    {
        printf '%s\n' "OK" >&1
    } 2>/dev/null || true
    exit 0
fi

# Function to decode MIME encoded subject (similar to the one used in extract_email_headers)
decode_mime_subject() {
    local raw_subject="$1"
    local decoded_subject="$raw_subject"
    
    # Check if subject is MIME encoded (contains =?...?= pattern)
    if [[ "$raw_subject" == *=?* ]]; then
        # Use Python for reliable MIME decoding
        if command -v python3 >/dev/null 2>&1; then
            local temp_py_script="/tmp/pmg_decode_subject_$$.py"
            cat > "$temp_py_script" << 'EOF_PY'
import sys
from email.header import decode_header
raw_subject = sys.argv[1] if len(sys.argv) > 1 else ""
decoded = decode_header(raw_subject)
result = []
for part, encoding in decoded:
    if isinstance(part, bytes):
        if encoding:
            result.append(part.decode(encoding))
        else:
            result.append(part.decode('utf-8'))
    else:
        result.append(part)
print(''.join(result))
EOF_PY
            if [ -f "$temp_py_script" ]; then
                decoded_subject=$(python3 "$temp_py_script" "$raw_subject" 2>/dev/null)
                if [ $? -ne 0 ] || [ -z "$decoded_subject" ]; then
                    decoded_subject="$raw_subject"
                fi
                rm -f "$temp_py_script"
            fi
        fi
    fi
    
    echo "$decoded_subject"
}

# Decode the subject for logging
decoded_subject=$(decode_mime_subject "$subject")

# Log email details before processing with decoded subject
log "Processing email - Sender: $sender, Receiver: $receiver, Subject: $decoded_subject"

# Limit email size to prevent API errors (e.g., 100KB limit)
EMAIL_SIZE_LIMIT=52400  # 100KB in bytes (100 * 1024)
email_size=$(echo -n "$email_content" | wc -c)

if [ "$email_size" -gt "$EMAIL_SIZE_LIMIT" ]; then
    log "Email too large ($email_size bytes), truncating to $EMAIL_SIZE_LIMIT bytes - Sender: $sender, Receiver: $receiver"
    # Truncate email content, handling potential pipe issues
    email_content=$(echo "$email_content" | head -c "$EMAIL_SIZE_LIMIT" 2>/dev/null || echo "${email_content:0:$EMAIL_SIZE_LIMIT}")
fi

# Prepare the request to OpenRouter API
request_body=$(cat <<EOF
{
  "model": "$OPENROUTER_MODEL",
  "messages": [
    {
      "role": "system",
      "content": "You are an email spam classifier. Analyze the following email and determine if it is spam. Respond with a JSON object containing 'is_spam' (boolean) and 'confidence' (number between 0 and 1) and 'reason' (string explaining why). Be strict in your assessment - only mark as spam if there are clear signs of unwanted or malicious content."
    },
    {
      "role": "user",
      "content": "Please analyze this email for spam content:\n\n$email_content"
    }
  ],
  "provider": {
    "sort": "price"
  }
}
EOF
)

# Log that we're about to make the API call
log "About to call OpenRouter API - Sender: $sender, Receiver: $receiver"

# Use cURL to call the OpenRouter API
response=$(curl -s -X POST \
  -H "Authorization: Bearer $OPENROUTER_API_KEY" \
  -H "Content-Type: application/json" \
  -d "$request_body" \
  "https://openrouter.ai/api/v1/chat/completions" \
  --max-time 240)  # 4 minute timeout to stay under the 5-minute limit

if [ $? -ne 0 ]; then
    log "Failed to get response from OpenRouter API - Sender: $sender, Receiver: $receiver"
    {
        printf '%s\n' "OK" >&1
    } 2>/dev/null || true
    exit 0
fi

if [ -z "$response" ]; then
    log "Empty response from OpenRouter API - Sender: $sender, Receiver: $receiver"
    {
        printf '%s\n' "OK" >&1
    } 2>/dev/null || true
    exit 0
fi

# Check if jq is available before using it
if ! command -v jq >/dev/null 2>&1; then
    log "ERROR: jq command not found - Sender: $sender, Receiver: $receiver"
    printf '%s\n' "OK" >&1
    exit 0
fi

response_size=$(echo -n "$response" | wc -c)
log "Successfully received response from OpenRouter API - Sender: $sender, Receiver: $receiver, Response size: $response_size bytes"

# Check if the response contains an error, but first check if jq is available
if command -v jq >/dev/null 2>&1; then
    if echo "$response" | jq -e '.error' >/dev/null 2>&1; then
        error_message=$(echo "$response" | jq -r '.error.message // "Unknown API error"')
        error_code=$(echo "$response" | jq -r '.error.code // "Unknown code"')
        log "OpenRouter API error - Sender: $sender, Receiver: $receiver, Error: $error_message, Code: $error_code"
        printf '%s\n' "OK" >&1
        exit 0
    fi
else
    log "ERROR: jq command not found during error checking - Sender: $sender, Receiver: $receiver"
    {
        printf '%s\n' "OK" >&1
    } 2>/dev/null || true
    exit 0
fi

# Extract the AI's response (need to extract the text content from the completion)
ai_response=$(echo "$response" | jq -r '.choices[0].message.content' 2>/dev/null)

if [ -z "$ai_response" ] || [ "$ai_response" = "null" ]; then
    log "Could not extract response from OpenRouter API - Sender: $sender, Receiver: $receiver, Response: $response"
    {
        printf '%s\n' "OK" >&1
    } 2>/dev/null || true
    exit 0
fi

# Check if sed is available before extracting JSON
if ! command -v sed >/dev/null 2>&1; then
    log "ERROR: sed command not found - Sender: $sender, Receiver: $receiver"
    {
        printf '%s\n' "OK" >&1
    } 2>/dev/null || true
    exit 0
fi

# Check if sed is available before extracting JSON
if ! command -v sed >/dev/null 2>&1; then
    log "ERROR: sed command not found - Sender: $sender, Receiver: $receiver"
    {
        printf '%s\n' "OK" >&1
    } 2>/dev/null || true
    exit 0
fi

# Try to extract JSON from the AI response (sometimes it's wrapped in markdown or additional text)
# Handle both markdown-wrapped JSON (```json ... ```) and plain JSON
json_response=""
if echo "$ai_response" | grep -q '```json'; then
    # Extract JSON from markdown code block
    json_response=$(echo "$ai_response" | sed -n '/```json/,/```/p' | head -n -1 | tail -n +2)
elif echo "$ai_response" | grep -E -q '\{.*\}'; then
    # Extract first JSON object found
    json_response=$(echo "$ai_response" | sed -n '/{/,/}/p' | sed -n 1p)
fi

# If extraction failed using above methods, try more flexible approach
if [ -z "$json_response" ] || [ "$json_response" = "null" ]; then
    # Find content between first { and last }
    if command -v grep >/dev/null 2>&1; then
        if echo "$ai_response" | grep -o '{.*}' >/dev/null 2>&1; then
            json_response=$(echo "$ai_response" | grep -o '{[^}]*}' | head -n 1)
        fi
    else
        log "ERROR: grep command not found during JSON extraction - Sender: $sender, Receiver: $receiver"
        {
            printf '%s\n' "OK" >&1
        } 2>/dev/null || true
        exit 0
    fi
fi

# Store the OpenRouter response for later use
openrouter_analysis="$ai_response"

# Log the complete OpenRouter analysis result
log "OpenRouter Analysis Result - Sender: $sender, Receiver: $receiver, Result: $ai_response"

# Detailed logging of the analysis results
if [ -n "$json_response" ] && [ "$json_response" != "null" ]; then
    # Check if jq is still available before parsing JSON response
    if command -v jq >/dev/null 2>&1; then
        is_spam=$(echo "$json_response" | jq -r '.is_spam' 2>/dev/null)
        confidence=$(echo "$json_response" | jq -r '.confidence' 2>/dev/null)
        reason=$(echo "$json_response" | jq -r '.reason' 2>/dev/null)
        
        log "Spam Analysis Result - Sender: $sender, Receiver: $receiver, Is Spam: $is_spam, Confidence: $confidence, Reason: $reason"
    else
        log "ERROR: jq command not found during JSON parsing - Sender: $sender, Receiver: $receiver"
        {
            printf '%s\n' "OK" >&1
        } 2>/dev/null || true
        exit 0
    fi
else
    log "Spam Analysis Result - Sender: $sender, Receiver: $receiver, Plain Text Analysis: $ai_response"
fi

if [ -n "$json_response" ] && [ "$json_response" != "null" ]; then
    # Parse the JSON response - check if jq is still available
    if command -v jq >/dev/null 2>&1; then
        is_spam=$(echo "$json_response" | jq -r '.is_spam' 2>/dev/null)
        confidence=$(echo "$json_response" | jq -r '.confidence' 2>/dev/null)
        reason=$(echo "$json_response" | jq -r '.reason' 2>/dev/null)
    else
        log "ERROR: jq command not found during JSON response parsing - Sender: $sender, Receiver: $receiver"
        printf '%s\n' "OK" >&1
        exit 0
    fi
    
    if [ "$is_spam" = "true" ]; then
        log "SPAM DETECTED - Sender: $sender, Receiver: $receiver, Confidence: $confidence, Reason: $reason, OpenRouter Analysis: $openrouter_analysis"
        
        # Calculate score from confidence
        score=$(awk "BEGIN {printf \"%.0f\", $confidence * 10}")
        
        # Extract domain from sender email address
        domain=$(echo "$sender" | grep -oE '@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}' | sed 's/^@//' | head -n 1)
        if [ -n "$domain" ]; then
            blacklist_file="$(dirname "$0")/blacklist.txt"
            # Check if domain is already in blacklist before adding
            if [ ! -f "$blacklist_file" ] || ! grep -Fxq "$domain" "$blacklist_file"; then
                log "Confirmed spam detected, adding domain $domain to blacklist"
                # Try to add to blacklist, but ignore errors
                {
                    echo "$domain" >> "$blacklist_file"
                } 2>/dev/null || true
            fi
        fi
        
        # Output the score result and exit immediately
        {
            printf '%s\n' "SCORE: $score" >&1
        } 2>/dev/null || true
        exit 0
    elif [ -n "$confidence" ] && [ "$confidence" != "null" ]; then
        # Check if awk is available for score calculation
        if ! command -v awk >/dev/null 2>&1; then
            log "ERROR: awk command not found - Sender: $sender, Receiver: $receiver"
            {
                printf '%s\n' "OK" >&1
            } 2>/dev/null || true
            exit 0
        fi
        
        # Convert confidence to a spam score between 0 and 10 using awk instead of bc
        score=$(awk "BEGIN {printf \"%.0f\", $confidence * 10}")
        log "Spam analysis completed - Sender: $sender, Receiver: $receiver, Confidence: $confidence, Score: $score, OpenRouter Analysis: $openrouter_analysis"

        # Only add score if it's meaningful (at least 0.3 confidence)
        if (( $(awk "BEGIN {print ($confidence > 0.3)}") )); then
            # If score is 7 or higher, add domain to blacklist
            if [ "$score" -ge 7 ]; then
                # Extract domain from sender email address
                domain=$(echo "$sender" | grep -oE '@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}' | sed 's/^@//' | head -n 1)
                if [ -n "$domain" ]; then
                    blacklist_file="$(dirname "$0")/blacklist.txt"
                    # Check if domain is already in blacklist before adding
                    if [ ! -f "$blacklist_file" ] || ! grep -Fxq "$domain" "$blacklist_file"; then
                        log "Score $score >= 7, adding domain $domain to blacklist"
                        # Try to add to blacklist, but ignore errors
                        {
                            echo "$domain" >> "$blacklist_file"
                        } 2>/dev/null || true
                    fi
                fi
            fi
            # Output the score result and exit immediately
            {
                printf '%s\n' "SCORE: $score" >&1
            } 2>/dev/null || true
            exit 0
        else
            # Output OK result and exit immediately
            {
                printf '%s\n' "OK" >&1
            } 2>/dev/null || true
            exit 0
        fi
    fi
else
    # If JSON parsing fails, try to infer from plain text response
    # Check if grep is available before using it
    if ! command -v grep >/dev/null 2>&1; then
        log "ERROR: grep command not found - Sender: $sender, Receiver: $receiver"
        {
            printf '%s\n' "OK" >&1
        } 2>/dev/null || true
        exit 0
    fi
    
    if echo "$ai_response" | grep -qi "spam\|phishing\|malware\|virus\|suspicious"; then
        log "SPAM DETECTED in plain text analysis - Sender: $sender, Receiver: $receiver, Analysis: $ai_response, OpenRouter Analysis: $openrouter_analysis"
        
        # Extract domain from sender email address
        domain=$(echo "$sender" | grep -oE '@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}' | sed 's/^@//' | head -n 1)
        if [ -n "$domain" ]; then
            blacklist_file="$(dirname "$0")/blacklist.txt"
            # Check if domain is already in blacklist before adding
            if [ ! -f "$blacklist_file" ] || ! grep -Fxq "$domain" "$blacklist_file"; then
                log "High confidence spam detected, adding domain $domain to blacklist"
                # Try to add to blacklist, but ignore errors
                {
                    echo "$domain" >> "$blacklist_file"
                } 2>/dev/null || true
            fi
        fi
        
        # Output the score result and exit immediately
        {
            printf '%s\n' "SCORE: 10" >&1
        } 2>/dev/null || true
        exit 0
    fi
fi

# Final log and result output - make sure nothing in this section fails
log "Email passed spam check - Sender: $sender, Receiver: $receiver, OpenRouter Analysis: $openrouter_analysis" 2>/dev/null || true
{
    printf '%s\n' "OK" >&1
} 2>/dev/null || true
exit 0

Дополнения
#

Скрипт должен находится по адресу /usr/local/bin/pmg-custom-check, в этой же папке whitelist.txt и blacklist.txt

А в /etc/pmg/pmg.conf должен быть включен функционал cutsom check

section: admin
    custom_check 1

И чтобы два раза не вставать, в PMG у меня уже есть какие-то домены в whitelist|blacklist, я решил их сразу добавить сюда. Сначала воспользуемся командой

sudo pmgsh get /config/ruledb/who

чтобы определить ID, созданных в PMG whitelist|blacklist, затем экспортируем домены во временные файлы

sudo pmgsh get /config/ruledb/who/2/objects | jq '.[] | select(.domain != null) | .domain' > /tmp/blacklist.text
sudo pmgsh get /config/ruledb/who/3/objects | jq '.[] | select(.domain != null) | .domain' > /tmp/whitelist.txt

После чего можем добавить их в /usr/local/bin/

Итог
#

Прикрутил AI к проверке почты. Я использовал OpenRouter и модель mistralai/mistral-small-3.1-24b-instruct (достаточно бюджетная). Сейчас я использую этот скрипт на не очень критичной почте, не уверен, что готов отправлять нейронке все свои сообщения. Можно использовать и какие-нибудь локальные llama-модели, но подходящего железа под рукой нет, а протестировать концепцию использования AI для анализа спама хотелось.