import json
from http.server import HTTPServer, BaseHTTPRequestHandler, ThreadingHTTPServer
import requests
import urllib3
import logging
import argparse

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
logging.getLogger("urllib3").setLevel(logging.WARNING)
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

class ProxyHandler(BaseHTTPRequestHandler):
    """
    自定义请求处理器，拦截所有 GET / POST / PUT / DELETE / OPTIONS 请求并转发
    """

    def do_request(self):
        target_url = f"{TARGET_BASE_URL}{self.path}"

        # 1. 过滤 Host，强制禁用压缩协商
        headers = {k: v for k, v in self.headers.items()
                   if k.lower() not in ['host', 'accept-encoding']}
        headers["Accept-Encoding"] = "identity"

        content_length = int(self.headers.get('Content-Length', 0))
        body = self.rfile.read(content_length) if content_length > 0 else None

        try:
            resp = requests.request(
                method=self.command,
                url=target_url,
                headers=headers,
                data=body,
                allow_redirects=False,
                verify=False
            )

            self.send_response(resp.status_code)

            # 跳过会冲突的头部
            skip_headers = {
                # 'transfer-encoding',
                # 'connection',
                'server',
                'date',
                # 'access-control-allow-origin',
                # 'access-control-allow-methods',
                # 'access-control-allow-headers',
                # 'access-control-allow-credentials',
                # 'access-control-max-age'
            }

            for key, value in resp.headers.items():
                if key.lower() not in skip_headers:
                    self.send_header(key, value)

            self.add_cors_headers()
            self.end_headers()
            self.wfile.write(resp.content)

        except Exception as e:
            logging.error(f"代理转发异常: {e}")
            error_msg = json.dumps({"error": str(e)})
            self.send_response(502)
            self.send_header("Content-Type", "application/json")
            self.add_cors_headers()
            self.end_headers()
            self.wfile.write(error_msg.encode('utf-8'))


    def add_cors_headers(self):
        """
        统一添加跨域资源共享 (CORS) 响应头
        """
        # 允许所有来源访问（如需限制特定域名，可将 '*' 替换为具体网址）
        self.send_header("Access-Control-Allow-Origin", "*")
        # 允许携带凭证（如 Cookie）
        self.send_header("Access-Control-Allow-Credentials", "true")
        # 允许前端使用的 HTTP 方法 (已添加 PUT 和 DELETE)
        self.send_header("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
        # 允许前端传递的请求头
        self.send_header("Access-Control-Allow-Headers", "*")

    def do_OPTIONS(self):
        """
        处理浏览器的跨域预检请求 (Preflight Request)
        直接返回 200 OK 及跨域头，无需向后端转发
        """
        self.send_response(200)
        self.add_cors_headers()
        self.end_headers()

    # 处理GET请求
    def do_GET(self):
        self.do_request()

    # 处理POST请求
    def do_POST(self):
        self.do_request()

    # 处理 PUT 请求
    def do_PUT(self):
        self.do_request()

    # 处理 DELETE 请求
    def do_DELETE(self):
        self.do_request()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(add_help=True)
    parser.add_argument('--host', type=str, default='127.0.0.1', help='监听ip')
    parser.add_argument('--port', type=int, default=54321, help='监听端口')
    parser.add_argument('--url', type=str, required=True, help='转发链接')
    args = parser.parse_args()

    TARGET_BASE_URL = args.url
    httpd = ThreadingHTTPServer((args.host, args.port), ProxyHandler)

    logging.warning(f"目标转发地址: {TARGET_BASE_URL}")
    logging.debug(f'代理服务器已启动 http://{args.host}:{args.port}')
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        logging.info("服务已停止")
        httpd.server_close()
   # python 1.py --host 0.0.0.0 --port 54321 --url https://www.baidu.com