python upload server

main.py

import http.server
import os
import re
import socketserver

PORT = 9000
UPLOAD_DIR = "./"


# 一个简单的辅助函数，用于解析 multipart/form-data 的头部
def parse_headers(headers_str):
    headers = {}
    for line in headers_str.split("\r\n"):
        if ":" in line:
            key, value = line.split(":", 1)
            headers[key.strip().lower()] = value.strip()
    return headers


class FileUploadHandler(http.server.BaseHTTPRequestHandler):
    def do_POST(self):
        try:
            # 1. 从 Content-Type 头中获取 boundary
            content_type = self.headers["Content-Type"]
            match = re.search(r"boundary=(.+)", content_type)
            if not match:
                self.send_error(400, "Content-Type header does not contain boundary")
                return
            boundary = "--" + match.group(1)

            # 2. 读取整个请求体
            content_length = int(self.headers["Content-Length"])
            body = self.rfile.read(content_length)

            # 3. 按 boundary 分割请求体
            parts = body.split(boundary.encode())

            # 第一个和最后一个部分通常是空的或只有结尾标记，我们忽略它们
            for part in parts[1:-1]:
                # 每个部分由头部和内容组成，用 \r\n\r\n 分隔
                if b"\r\n\r\n" not in part:
                    continue

                header_bytes, content = part.split(b"\r\n\r\n", 1)
                # 去掉开头的 \r\n
                header_bytes = header_bytes.strip()

                # 解码头部信息
                headers_str = header_bytes.decode("utf-8")
                part_headers = parse_headers(headers_str)

                # 4. 解析 Content-Disposition 获取文件名
                if "content-disposition" not in part_headers:
                    continue

                disposition = part_headers["content-disposition"]
                # 使用正则表达式从 "form-data; name="file"; filename="example.txt"" 中提取文件名
                filename_match = re.search(r'filename="([^"]+)"', disposition)

                if not filename_match:
                    continue

                filename = filename_match.group(1)
                # 基本的安全过滤，防止路径遍历
                filename = os.path.basename(filename)

                if filename:
                    # 确保上传目录存在
                    os.makedirs(UPLOAD_DIR, exist_ok=True)

                    # 5. 将文件内容写入本地
                    filepath = os.path.join(UPLOAD_DIR, filename)
                    # 去掉最后的 \r\n
                    file_content = content.rstrip(b"\r\n")

                    with open(filepath, "wb") as f:
                        f.write(file_content)

                    print(f"Successfully received and saved file: {filename}")

            # 发送成功响应
            self.send_response(200)
            self.end_headers()
            self.wfile.write(b"File(s) uploaded successfully.")

        except Exception as e:
            print(f"An error occurred: {e}")
            self.send_error(500, f"Server error: {e}")


# --- 启动服务器 ---
Handler = FileUploadHandler

with socketserver.TCPServer(("", PORT), Handler) as httpd:
    print(f"Serving at port {PORT}")
    print(f"Upload files to http://localhost:{PORT}/ using `curl -X POST http://127.0.0.1:9000/ -F "file=@xxx/yyy.log"`")
    httpd.serve_forever()