shirobana-bot/mail-bridge/admin_com.py

"""Authenticated, deliberately limited hMailServer operations for the remote bot."""
import base64
import binascii
import hmac
import logging
import os
import re
from flask import Blueprint, jsonify, request
import com_backend as backend

log = logging.getLogger(__name__)
bp = Blueprint('admin_com', __name__, url_prefix='/v1/mail/api/admin/com')


def _positive_env_int(name, default):
    try:
        value = int(os.environ.get(name, default))
    except (TypeError, ValueError):
        return default
    return value if value > 0 else default


MAX_REQUEST_BYTES = 1024 * 1024
ATTACHMENT_REQUEST_MAX_BYTES = _positive_env_int(
    'MAIL_ATTACHMENT_REQUEST_MAX_BYTES', 40 * 1024 * 1024)
MAX_ATTACHMENT_COUNT = _positive_env_int('MAIL_ATTACHMENT_MAX_COUNT', 10)
MAX_ATTACHMENT_BYTES = _positive_env_int(
    'MAIL_ATTACHMENT_MAX_FILE_BYTES', 20 * 1024 * 1024)
MAX_ATTACHMENT_TOTAL_BYTES = _positive_env_int(
    'MAIL_ATTACHMENT_MAX_TOTAL_BYTES', 25 * 1024 * 1024)


@bp.before_request
def authenticate():
    token = backend.CONFIG.get('token', '')
    supplied = request.headers.get('Authorization', '')
    if len(token) < 64 or not hmac.compare_digest(supplied.encode(), ('Bearer '+token).encode()):
        return jsonify(error='unauthorized'), 401
    is_attachment_send = request.endpoint == 'admin_com.send_with_attachments'
    limit = ATTACHMENT_REQUEST_MAX_BYTES if is_attachment_send else MAX_REQUEST_BYTES
    # 新端点拒绝无 Content-Length 的流式请求，避免绕过入口大小限制；
    # 其他端点保留原有 1 MiB 行为和响应格式。
    if is_attachment_send and request.content_length is None:
        return jsonify(error='content length required'), 411
    if request.content_length and request.content_length > limit:
        return jsonify(error='request too large'), 413


@bp.after_request
def no_cache(response):
    response.headers['Cache-Control'] = 'no-store'
    return response


def address(value, local=True):
    if not isinstance(value, str) or not re.fullmatch(r'[^\s<>@,;]+@[^\s<>@,;]+', value):
        raise ValueError('invalid address')
    value = value.lower()
    if local and value.rsplit('@', 1)[1] not in backend.CONFIG['allowed_domains']:
        raise ValueError('unsupported domain')
    return value


def integer(value, low, high):
    if isinstance(value, bool): raise ValueError('invalid integer')
    value = int(value)
    if value < low or value > high: raise ValueError('integer out of range')
    return value


def attachment_filename(value):
    """只接受 MIME 显示文件名，绝不接受客户端提供的服务器路径。"""
    if not isinstance(value, str) or not value or len(value) > 255:
        raise ValueError('invalid attachment filename')
    if value != value.strip() or value in ('.', '..'):
        raise ValueError('invalid attachment filename')
    if len(value.encode('utf-8')) > 1024:
        raise ValueError('invalid attachment filename')
    if any(ord(ch) < 32 or ord(ch) == 127 or ch in '<>:"/\\|?*' for ch in value):
        raise ValueError('invalid attachment filename')
    return value


def attachment_content_type(value):
    # 仅 type/subtype，不含参数；同时排除 CR/LF 与其他头部注入字符。
    if not isinstance(value, str) or len(value) > 127:
        raise ValueError('invalid attachment MIME type')
    if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9!#$&^_.+-]*/"
                        r"[A-Za-z0-9][A-Za-z0-9!#$&^_.+-]*", value):
        raise ValueError('invalid attachment MIME type')
    return value.lower()


def incoming_attachments(value):
    if not isinstance(value, list) or not 1 <= len(value) <= MAX_ATTACHMENT_COUNT:
        raise ValueError('invalid attachment count')
    result, total = [], 0
    max_encoded_size = 4 * ((MAX_ATTACHMENT_BYTES + 2) // 3)
    for item in value:
        # 白名单字段可防止把 path / filename_on_server 等字段带入后端。
        if not isinstance(item, dict) or set(item) != {'filename', 'content_type', 'data_base64'}:
            raise ValueError('invalid attachment')
        filename = attachment_filename(item['filename'])
        content_type = attachment_content_type(item['content_type'])
        encoded = item['data_base64']
        if not isinstance(encoded, str) or len(encoded) > max_encoded_size:
            raise ValueError('attachment too large')
        try:
            content = base64.b64decode(encoded.encode('ascii'), validate=True)
        except (UnicodeError, binascii.Error, ValueError):
            raise ValueError('invalid attachment base64')
        if len(content) > MAX_ATTACHMENT_BYTES:
            raise ValueError('attachment too large')
        total += len(content)
        if total > MAX_ATTACHMENT_TOTAL_BYTES:
            raise ValueError('attachments too large')
        result.append({'filename': filename, 'content_type': content_type,
                       'content': content})
    return result, total


def item(m):
    return {'id': int(m.ID), 'subject': (m.Subject or '').strip() or '(无主题)',
            'from': (m.From or '').strip(), 'from_address': (m.FromAddress or '').strip(),
            'ts': backend._message_ts(m)}


@bp.route('/status', methods=['GET'])
def status():
    with backend._ComSession() as app:
        return jsonify(ok=True, version=app.Version)


@bp.route('/inbox', methods=['POST'])
def inbox():
    data = request.get_json()
    addr = address(data.get('address'))
    limit = integer(data.get('limit', 10), 1, 100)
    window = data.get('window_seconds')
    if window is not None: window = integer(window, 0, 315360000)
    with backend._ComSession() as app:
        result = backend._list_inbox_messages(app, addr, window, limit)
    return jsonify(items=result)


@bp.route('/message', methods=['POST'])
def message():
    data = request.get_json()
    addresses = data.get('addresses')
    if not isinstance(addresses, list) or not 1 <= len(addresses) <= 30:
        raise ValueError('invalid address list')
    addresses = [address(a) for a in addresses]
    mid = integer(data.get('id'), 1, 2**63-1)
    with backend._ComSession() as app:
        m = backend._find_message_by_id(app, addresses, mid)
        if m is None: return jsonify(message=None)
        images, attachments = backend._analyze_message(m)
        result = dict(ID=int(m.ID), Subject=m.Subject or '', From=m.From or '',
                      FromAddress=m.FromAddress or '', Date=str(m.Date or ''),
                      Body=m.Body or '', HTMLBody=m.HTMLBody or '',
                      image_count=images, attachment_names=attachments)
    return jsonify(message=result)


@bp.route('/changes', methods=['POST'])
def changes():
    data = request.get_json()
    addr = address(data.get('address'))
    after = data.get('after')
    if after is not None: after = integer(after, 0, 2**63-1)
    limit = integer(data.get('limit', 10), 1, 100)
    with backend._ComSession() as app:
        msgs = backend._get_account(app, addr).IMAPFolders.ItemByName('INBOX').Messages
        found, max_id = [], 0
        for idx in range(msgs.Count):
            m = msgs[idx]
            mid = int(m.ID)
            max_id = max(max_id, mid)
            if after is not None and mid > after: found.append(item(m))
        found.sort(key=lambda v: v['id'])
    return jsonify(max_id=max_id, items=found[:limit])


def outgoing(data=None):
    if data is None:
        data = request.get_json()
    sender, recipient = address(data.get('from_address')), address(data.get('to_address'), local=False)
    subject, body = data.get('subject'), data.get('body_text')
    if not isinstance(subject, str) or len(subject) > 2000 or '\r' in subject or '\n' in subject:
        raise ValueError('invalid subject')
    if not isinstance(body, str) or len(body) > 200000: raise ValueError('invalid body')
    return sender, recipient, subject, body


@bp.route('/send', methods=['POST'])
def send():
    args = outgoing()
    with backend._ComSession() as app:
        backend._get_account(app, args[0])
        backend._send_mail(app, *args)
    return jsonify(accepted=True)


@bp.route('/send_with_attachments', methods=['POST'])
def send_with_attachments():
    data = request.get_json()
    if not isinstance(data, dict):
        raise ValueError('invalid request')
    args = outgoing(data)
    if not args[2].strip():
        raise ValueError('invalid subject')
    attachments, total = incoming_attachments(data.get('attachments'))
    with backend._ComSession() as app:
        backend._get_account(app, args[0])
        sent_saved = backend._send_mail_with_attachments(app, *args, attachments)
    return jsonify(ok=True, accepted=True, sent_saved=sent_saved,
                   attachment_names=[a['filename'] for a in attachments],
                   size=total)


@bp.route('/sent', methods=['POST'])
def sent():
    args = outgoing()
    with backend._ComSession() as app:
        backend._save_to_sent(app, *args)
    return jsonify(ok=True)


@bp.errorhandler(Exception)
def error(exc):
    from werkzeug.exceptions import HTTPException
    if isinstance(exc, HTTPException): return jsonify(error=exc.name), exc.code
    if isinstance(exc, (ValueError, TypeError, AttributeError)):
        return jsonify(error='invalid request'), 400
    log.exception('COM operation failed')
    return jsonify(error='mail backend unavailable'), 502


def install(app):
    app.register_blueprint(bp)