floorplan-quotation/secure_preview_server.py

"""Local static server and credential-safe Volcengine Ark image proxy."""

import json
import os
import re
import base64
import binascii
import urllib.error
import urllib.request
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer


PORT = int(os.environ.get("PORT", "8000"))
MAX_BODY_BYTES = 16 * 1024 * 1024
ARK_IMAGE_URL = "https://ark.cn-beijing.volces.com/api/v3/images/generations"
ARK_CHAT_URL = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
ARK_VISION_MODEL = os.environ.get("ARK_VISION_MODEL", "doubao-1.5-vision-pro-32k-250115").strip()
RUNTIME_ARK_API_KEY = os.environ.get("ARK_API_KEY", "").strip()
EDIT_PROMPT = (
    "严格保持参考户型图的外轮廓、承重墙、隔墙、门洞、窗洞的位置、比例和数量完全不变。"
    "删除全部中英文文字、尺寸数字、红蓝标注线、箭头、家具、柜体、装饰图案、面积标注、logo和背景水印。"
    "不要增加、移动、弯曲或删减任何墙体、门洞和窗洞。输出完整、居中、无裁切的专业建筑平面底图："
    "纯白背景，墙体使用清晰深灰色粗线，门窗使用简洁细线，房间内部留白；不得包含文字、数字、家具、标注、logo或水印。"
)


DATA_URL_RE = re.compile(r"^data:image/([A-Za-z0-9.+-]+);base64,(.+)$", re.S)


def _normalize_image(value):
    """校验并规范化前端传来的图片：必须是标准 data URI，base64 合法且不超过10MB。

    去除空白/换行、小写化MIME、jpg→jpeg，确保转发给方舟的参数永远是合法格式，
    从根源上避免模型侧报 invalid_parameter_error（URL 无效）。
    """
    image = str(value or "").strip()
    image = re.sub(r"\s+", "", image)
    match = DATA_URL_RE.match(image)
    if not match:
        raise ValueError("图片必须是标准 Base64 Data URL（data:image/<格式>;base64,...），请重新上传户型图")
    fmt = match.group(1).lower()
    if fmt == "jpg":
        fmt = "jpeg"
    if fmt not in ("jpeg", "png", "webp", "bmp", "tiff", "gif"):
        raise ValueError("图片格式不支持：" + fmt)
    payload = match.group(2)
    try:
        raw = base64.b64decode(payload, validate=True)
    except (ValueError, binascii.Error):
        raise ValueError("Base64 图片内容无效，请重新上传户型图")
    if len(raw) > 10 * 1024 * 1024:
        raise ValueError("图片超过10MB限制，请压缩后重新上传")
    return "data:image/" + fmt + ";base64," + payload


class Handler(SimpleHTTPRequestHandler):
    # HTTP/1.1 keep-alive：避免 HTTP/1.0 每请求关连接，预览面板复用连接时遇到 RST 而 ERR_ABORTED
    protocol_version = "HTTP/1.1"

    def _vite(self):
        # Trae 预览会注入 /@vite/client 热更新脚本：返回 200 空JS，避免 404 中断主文档加载
        self.send_response(200)
        self.send_header("Content-Type", "text/javascript")
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", "2")
        self.end_headers()
        self.wfile.write(b"//")

    def _strip_cond(self):
        # 剥离条件请求头 → 永远 200 全量正文：
        # 预览面板需改写 HTML 正文注入热更新脚本，304 无正文会使其放弃本次加载并重发（net::ERR_ABORTED）
        for h in ("If-Modified-Since", "If-None-Match"):
            if h in self.headers:
                del self.headers[h]

    def send_header(self, keyword, value):
        # Python 3.14 的 send_head 用 "Content-type"（小写t），比较必须忽略大小写
        if keyword.lower() == "content-type" and "text/html" in value:
            super().send_header("Cache-Control", "no-store")
        super().send_header(keyword, value)

    def do_GET(self):
        # 【加固】静态文件服务遇到浏览器中断连接（预览重载/取消）时静默忽略，
        # 避免异常冒泡导致线程 traceback 刷屏甚至进程退出
        try:
            if self.path.split("?", 1)[0].startswith("/@vite/"):
                self._vite()
                return
            self._strip_cond()
            super().do_GET()
        except (BrokenPipeError, ConnectionResetError, TimeoutError):
            pass

    def do_HEAD(self):
        try:
            if self.path.split("?", 1)[0].startswith("/@vite/"):
                self._vite()
                return
            self._strip_cond()
            super().do_HEAD()
        except (BrokenPipeError, ConnectionResetError, TimeoutError):
            pass

    def _json_error(self, status, message):
        body = json.dumps({"code": "LocalProxyError", "message": message}, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _read_body(self, limit=MAX_BODY_BYTES):
        """无论后续成功或报错，都必须先把请求体完整读完再响应。

        否则提前返回并关闭连接时，套接字里残留的未读请求数据会触发 TCP RST，
        浏览器端表现为 net::ERR_CONNECTION_RESET（尤其是上传大图片 base64 时）。
        """
        try:
            length = int(self.headers.get("Content-Length", "0") or "0")
        except ValueError:
            length = 0
        if length <= 0:
            if length < 0:
                raise ValueError("请求内容无效")
            return b""
        oversized = length > limit
        received = 0
        chunks = []
        while received < length:
            chunk = self.rfile.read(min(length - received, 262144))
            if not chunk:
                break
            received += len(chunk)
            if not oversized:
                chunks.append(chunk)
        if received < length:
            raise ValueError("上传内容不完整")
        if oversized:
            raise ValueError("上传内容超过大小限制，请压缩后重试")
        return b"".join(chunks)

    def do_POST(self):
        global RUNTIME_ARK_API_KEY
        route = self.path.split("?", 1)[0]
        if route == "/api/export/excel":
            try:
                raw = self._read_body(MAX_BODY_BYTES)
                if not raw:
                    raise ValueError("Excel文件为空")
                payload = json.loads(raw.decode("utf-8"))
                filename = os.path.basename(str(payload.get("filename", "报价单.xlsx"))).strip()
                if not filename.lower().endswith(".xlsx"):
                    filename += ".xlsx"
                filename = "".join(c for c in filename if c not in '<>:"/\\|?*') or "报价单.xlsx"
                binary = base64.b64decode(str(payload.get("data", "")), validate=True)
                if not binary:
                    raise ValueError("Excel文件内容为空")
                # 允许前端指定绝对目录（仅限本机盘符路径），缺省回退桌面
                dir_param = str(payload.get("dir", "")).strip()
                if dir_param and re.match(r"^[A-Za-z]:[\\/]", dir_param) and os.path.isdir(dir_param):
                    desktop = dir_param
                else:
                    desktop = os.path.join(os.path.expanduser("~"), "Desktop")
                os.makedirs(desktop, exist_ok=True)
                target = os.path.join(desktop, filename)
                stem, ext = os.path.splitext(target)
                counter = 2
                while os.path.exists(target):
                    target = f"{stem}_{counter}{ext}"
                    counter += 1
                with open(target, "wb") as output:
                    output.write(binary)
                result = json.dumps({"saved": True, "filename": os.path.basename(target), "path": target}, ensure_ascii=False).encode("utf-8")
            except (ValueError, json.JSONDecodeError, UnicodeDecodeError, base64.binascii.Error) as error:
                self._json_error(400, str(error))
                return
            self.send_response(200)
            self.send_header("Content-Type", "application/json; charset=utf-8")
            self.send_header("Cache-Control", "no-store")
            self.send_header("Content-Length", str(len(result)))
            self.end_headers()
            self.wfile.write(result)
            return
        if route == "/api/ark/config":
            try:
                raw = self._read_body(4096)
                if not raw:
                    raise ValueError("请求内容无效")
                payload = json.loads(raw.decode("utf-8"))
                key = str(payload.get("api_key", "")).strip()
                if not key.startswith("ark-") or len(key) < 20 or len(key) > 256:
                    raise ValueError("模型 Key 格式不正确")
                RUNTIME_ARK_API_KEY = key
            except (ValueError, json.JSONDecodeError, UnicodeDecodeError) as error:
                self._json_error(400, str(error))
                return
            result = json.dumps({"configured": True}, ensure_ascii=False).encode("utf-8")
            self.send_response(200)
            self.send_header("Content-Type", "application/json; charset=utf-8")
            self.send_header("Cache-Control", "no-store")
            self.send_header("Content-Length", str(len(result)))
            self.end_headers()
            self.wfile.write(result)
            return
        if route == "/api/ark/dimension-detect":
            # 先完整读取并校验请求体（含图片规范化），再检查 Key，
            # 保证任何分支都不会在未读完请求体时提前关闭连接
            try:
                raw = self._read_body(MAX_BODY_BYTES)
                if not raw:
                    raise ValueError("上传内容为空")
                incoming = json.loads(raw.decode("utf-8"))
                image = _normalize_image(incoming.get("image", ""))
            except (ValueError, json.JSONDecodeError, UnicodeDecodeError) as error:
                self._json_error(400, str(error))
                return
            api_key = RUNTIME_ARK_API_KEY
            if not api_key:
                self._json_error(503, "本机后端尚未配置 ARK_API_KEY")
                return
            prompt = (
                "你是建筑户型图尺寸标注识别器。检查图片中已经存在的线性真实尺寸标注，"
                "识别尺寸数字及其对应的双箭头、箭头或尺寸线两端。忽略面积数值、房间名称、柜体编号和计算说明。"
                "只输出JSON，不要Markdown：{\"dimensions\":[{\"value\":3150,\"unit\":\"mm\","
                "\"x1\":5.2,\"y1\":42.1,\"x2\":5.2,\"y2\":68.4,\"confidence\":0.96}]}."
                "x1,y1,x2,y2是箭头两端相对于整张图片宽高的0到100百分比坐标。"
                "仅返回有明确尺寸数字且能定位对应线段两端的项目；没有则返回{\"dimensions\":[]}。"
            )
            body = json.dumps({
                "model": ARK_VISION_MODEL,
                "messages": [{"role": "user", "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": image}}
                ]}],
                "stream": False,
                "max_tokens": 1600,
                "temperature": 0.1,
            }, ensure_ascii=False).encode("utf-8")
            request = urllib.request.Request(ARK_CHAT_URL, data=body, method="POST", headers={
                "Authorization": "Bearer " + api_key,
                "Content-Type": "application/json",
            })
            try:
                with urllib.request.urlopen(request, timeout=120) as response:
                    result, status = response.read(), response.status
            except urllib.error.HTTPError as error:
                result, status = error.read(), error.code
            except Exception as error:
                self._json_error(502, "连接火山方舟视觉模型失败：" + str(error))
                return
            self.send_response(status)
            self.send_header("Content-Type", "application/json; charset=utf-8")
            self.send_header("Cache-Control", "no-store")
            self.send_header("Content-Length", str(len(result)))
            self.end_headers()
            self.wfile.write(result)
            return
        if route != "/api/ark/image-edit":
            try:
                self._read_body(MAX_BODY_BYTES)
            except ValueError:
                pass
            self._json_error(404, "接口不存在")
            return
        # 同样先读完请求体并校验，再检查 Key，避免提前关闭连接导致 ERR_CONNECTION_RESET
        try:
            raw = self._read_body(MAX_BODY_BYTES)
            if not raw:
                raise ValueError("上传内容为空")
            incoming = json.loads(raw.decode("utf-8"))
            image = _normalize_image(incoming.get("image", ""))
        except (ValueError, json.JSONDecodeError, UnicodeDecodeError) as error:
            self._json_error(400, str(error))
            return
        api_key = RUNTIME_ARK_API_KEY
        if not api_key:
            self._json_error(503, "本机后端尚未配置 ARK_API_KEY")
            return

        body = json.dumps({
            "model": "doubao-seedream-4-5-251128",
            "prompt": EDIT_PROMPT,
            "image": image,
            "size": "2K",
            "sequential_image_generation": "disabled",
            "stream": False,
            "response_format": "b64_json",
            "watermark": False,
            "guidance_scale": 8.5,
        }, ensure_ascii=False).encode("utf-8")
        request = urllib.request.Request(
            ARK_IMAGE_URL,
            data=body,
            method="POST",
            headers={
                "Authorization": "Bearer " + api_key,
                "Content-Type": "application/json",
            },
        )
        try:
            with urllib.request.urlopen(request, timeout=180) as response:
                result, status = response.read(), response.status
        except urllib.error.HTTPError as error:
            result, status = error.read(), error.code
        except Exception as error:
            self._json_error(502, "连接火山方舟失败：" + str(error))
            return
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", str(len(result)))
        self.end_headers()
        self.wfile.write(result)


class SafeThreadingHTTPServer(ThreadingHTTPServer):
    # 【加固】单个请求线程异常只记日志不杀进程，保证预览服务长期在线
    def handle_error(self, request, client_address):
        import sys
        print(f"[server] 请求处理异常（已忽略，服务继续）: {client_address}", file=sys.stderr)


if __name__ == "__main__":
    print(f"Secure preview: http://localhost:{PORT}/AI快速报价_Demo2.html")
    if not os.environ.get("ARK_API_KEY"):
        print("Warning: ARK_API_KEY is not configured; Doubao image editing is disabled.")
    SafeThreadingHTTPServer(("127.0.0.1", PORT), Handler).serve_forever()