client.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
HTTP客户端
发送HTTP请求，支持指定case-id header
"""

import json
import requests
import argparse
import os
from typing import Optional, Dict, Any
from case_manager import CaseManager


class HTTPClient:
    """HTTP客户端类"""
    
    def __init__(self, base_url: str = 'http://localhost:32941'):
        self.base_url = base_url
        self.case_id_header = 'X-Case-ID'
    
    def send_request(
        self,
        case_id: str,
        method: str = 'GET',
        path: str = '/',
        headers: Optional[Dict[str, str]] = None,
        data: Optional[Any] = None,
        json_data: Optional[Dict] = None
    ) -> requests.Response:
        """
        发送HTTP请求
        
        Args:
            case_id: 用例ID，会添加到header中
            method: HTTP方法 (GET, POST, PUT, DELETE等)
            path: 请求路径
            headers: 额外的请求头
            data: 请求体（字符串或字节）
            json_data: JSON格式的请求体
        
        Returns:
            requests.Response对象
        """
        url = f"{self.base_url.rstrip('/')}/{path.lstrip('/')}"
        
        # 构建请求头
        request_headers = headers.copy() if headers else {}
        request_headers[self.case_id_header] = case_id
        
        # 发送请求
        response = requests.request(
            method=method.upper(),
            url=url,
            headers=request_headers,
            data=data,
            json=json_data,
            timeout=30
        )
        
        return response
    
    def send_from_case(self, case_id: str, request_data: Dict[str, Any]) -> requests.Response:
        """
        根据用例中的请求数据发送请求
        
        Args:
            case_id: 用例ID
            request_data: 包含请求信息的字典，格式：
                {
                    'method': 'GET',
                    'path': '/api/users',
                    'headers': {...},
                    'body': '...' 或 {...}
                }
        
        Returns:
            requests.Response对象
        """
        method = request_data.get('method', 'GET')
        path = request_data.get('path', '/')
        headers = request_data.get('headers', {})
        body = request_data.get('body', None)
        
        # 处理body
        json_data = None
        data = None
        
        if body:
            if isinstance(body, dict):
                json_data = body
            elif isinstance(body, str):
                # 尝试解析JSON字符串
                try:
                    json_data = json.loads(body)
                except json.JSONDecodeError:
                    data = body
        
        return self.send_request(
            case_id=case_id,
            method=method,
            path=path,
            headers=headers,
            data=data,
            json_data=json_data
        )


def main():
    """命令行入口"""
    parser = argparse.ArgumentParser(description='HTTP客户端工具')
    parser.add_argument('case_id', nargs='?', help='用例ID（可选，不指定时重放所有用例）')
    parser.add_argument('--method', '-m', help='HTTP方法（仅当指定单个用例时有效）')
    parser.add_argument('--path', '-p', help='请求路径（仅当指定单个用例时有效）')
    parser.add_argument('--header', '-H', action='append', help='请求头，格式: Key:Value（仅当指定单个用例时有效）')
    parser.add_argument('--data', '-d', help='请求体数据（仅当指定单个用例时有效）')
    parser.add_argument('--json', '-j', help='JSON格式的请求体（仅当指定单个用例时有效）')
    parser.add_argument('--url', '-u', help='服务器完整URL (例如: http://localhost:23941)')
    parser.add_argument('--host', default='localhost', help='服务器IP地址 (默认: localhost)')
    parser.add_argument('--port', type=int, default=23941, help='服务器端口 (默认: 23941)')
    parser.add_argument('--pretty', action='store_true', help='美化输出JSON响应')
    parser.add_argument('--cases', '-c', help='用例文件或文件夹路径（用于加载所有用例，默认: cases.yaml和cases/）')
    
    args = parser.parse_args()
    
    # 构建服务器URL
    if args.url:
        # 如果指定了完整URL，使用它
        server_url = args.url
    else:
        # 否则使用host和port构建URL
        server_url = f'http://{args.host}:{args.port}'
    
    # 创建客户端
    client = HTTPClient(base_url=server_url)
    
    # 如果没有指定case_id，重放所有用例
    if not args.case_id:
        return replay_all_cases(client, server_url, args)
    
    # 指定了单个用例，执行单个请求
    return send_single_request(client, server_url, args)


def send_single_request(client, server_url, args):
    """发送单个用例请求"""
    # 解析请求头
    headers = {}
    if args.header:
        for h in args.header:
            if ':' in h:
                key, value = h.split(':', 1)
                headers[key.strip()] = value.strip()
    
    # 解析JSON数据
    json_data = None
    if args.json:
        try:
            json_data = json.loads(args.json)
        except json.JSONDecodeError as e:
            print(f"Error parsing JSON: {e}")
            return 1
    
    # 使用命令行参数或从用例文件加载请求信息
    method = args.method or 'GET'
    path = args.path or '/'
    
    # 尝试从用例文件加载请求信息
    case_manager = CaseManager()
    if args.cases:
        case_manager.load_cases(args.cases)
    else:
        # 默认加载cases.yaml和cases/文件夹
        case_manager.load_cases('cases.yaml')
        if os.path.exists('cases') and os.path.isdir('cases'):
            case_manager.load_cases('cases')
    
    case = case_manager.get_case(args.case_id)
    if case and case.get('request'):
        # 使用用例文件中的请求信息
        request_data = case['request']
        method = request_data.get('method', method)
        path = request_data.get('path', path)
        # 合并用例中的headers和命令行指定的headers
        case_headers = request_data.get('headers', {})
        case_headers.update(headers)
        headers = case_headers
        
        # 使用用例中的body
        case_body = request_data.get('body')
        if case_body is not None:
            if isinstance(case_body, dict):
                json_data = case_body
            elif isinstance(case_body, str):
                try:
                    json_data = json.loads(case_body)
                except json.JSONDecodeError:
                    json_data = None
                    args.data = case_body
    
    # 构建完整的请求信息
    url = f"{server_url.rstrip('/')}/{path.lstrip('/')}"
    request_headers = headers.copy() if headers else {}
    request_headers[client.case_id_header] = args.case_id
    
    # 解析URL获取host和path
    from urllib.parse import urlparse
    parsed_url = urlparse(server_url)
    host = parsed_url.netloc or 'localhost'
    path = path if path.startswith('/') else '/' + path
    
    # 构建请求体
    request_body_str = None
    if json_data:
        request_body_str = json.dumps(json_data, ensure_ascii=False) if not args.pretty else json.dumps(json_data, indent=2, ensure_ascii=False)
    elif args.data:
        request_body_str = str(args.data)
    
    # 获取用例描述
    case = case_manager.get_case(args.case_id)
    case_description = case_manager.get_case_description(args.case_id) if case else args.case_id
    
    # 显示请求信息（BurpSuite风格）
    print("\n" + "=" * 80)
    print(f"REQUEST [{case_description}]")
    print("=" * 80)
    
    # 请求行
    print(f"{method.upper()} {path} HTTP/1.1")
    
    # 请求头（按BurpSuite顺序）
    # 先显示Host
    if 'Host' in request_headers:
        print(f"Host: {request_headers['Host']}")
    else:
        print(f"Host: {host}")
    
    # 显示其他请求头
    for key, value in request_headers.items():
        if key.lower() != 'host':
            print(f"{key}: {value}")
    
    # 空行分隔header和body
    print()
    
    # 请求体
    if request_body_str:
        print(request_body_str)
    
    print("=" * 80)
    print("RESPONSE")
    print("=" * 80)
    
    try:
        # 发送请求
        response = client.send_request(
            case_id=args.case_id,
            method=method,
            path=path,
            headers=headers,
            data=args.data,
            json_data=json_data
        )
        
        # 状态行
        status_text_map = {
            200: 'OK', 201: 'Created', 202: 'Accepted', 204: 'No Content',
            301: 'Moved Permanently', 302: 'Found', 304: 'Not Modified',
            400: 'Bad Request', 401: 'Unauthorized', 403: 'Forbidden',
            404: 'Not Found', 405: 'Method Not Allowed', 500: 'Internal Server Error',
            502: 'Bad Gateway', 503: 'Service Unavailable'
        }
        status_text = status_text_map.get(response.status_code, 'OK' if 200 <= response.status_code < 300 else 'Error')
        print(f"HTTP/1.1 {response.status_code} {status_text}")
        
        # 响应头
        for key, value in response.headers.items():
            print(f"{key}: {value}")
        
        # 空行分隔header和body
        print()
        
        # 响应体
        try:
            response_json = response.json()
            if args.pretty:
                print(json.dumps(response_json, indent=2, ensure_ascii=False))
            else:
                print(json.dumps(response_json, ensure_ascii=False))
        except (ValueError, json.JSONDecodeError):
            # 不是JSON格式，直接输出文本
            print(response.text)
        
        print("=" * 80 + "\n")
        return 0
        
    except requests.exceptions.RequestException as e:
        print(f"HTTP/1.1 000 Connection Failed")
        print()
        print(f"Request failed: {e}")
        print("=" * 80 + "\n")
        return 1


def replay_all_cases(client, server_url, args):
    """重放所有用例"""
    # 加载用例
    case_manager = CaseManager()
    if args.cases:
        case_manager.load_cases(args.cases)
    else:
        # 默认加载cases.yaml和cases/文件夹
        case_manager.load_cases('cases.yaml')
        if os.path.exists('cases') and os.path.isdir('cases'):
            case_manager.load_cases('cases')
    
    case_ids = case_manager.list_cases()
    
    if not case_ids:
        print("No cases found. Please check your case files.")
        return 1
    
    print(f"\nFound {len(case_ids)} case(s). Starting replay...\n")
    print("=" * 80)
    
    success_count = 0
    fail_count = 0
    
    for case_id in sorted(case_ids):
        case = case_manager.get_case(case_id)
        if not case or not case.get('request'):
            case_description = case_manager.get_case_description(case_id)
            print(f"\n[SKIP] {case_description}: No request data found")
            continue
        
        request_data = case['request']
        method = request_data.get('method', 'GET')
        path = request_data.get('path', '/')
        headers = request_data.get('headers', {})
        body = request_data.get('body')
        
        # 获取用例描述
        case_description = case_manager.get_case_description(case_id)
        
        # 准备请求数据
        json_data = None
        data = None
        if body is not None:
            if isinstance(body, dict):
                json_data = body
            elif isinstance(body, str):
                try:
                    json_data = json.loads(body)
                except json.JSONDecodeError:
                    data = body
        
        # 发送请求
        try:
            response = client.send_request(
                case_id=case_id,
                method=method,
                path=path,
                headers=headers,
                data=data,
                json_data=json_data
            )
            
            # 显示结果
            status_text_map = {
                200: 'OK', 201: 'Created', 202: 'Accepted', 204: 'No Content',
                301: 'Moved Permanently', 302: 'Found', 304: 'Not Modified',
                400: 'Bad Request', 401: 'Unauthorized', 403: 'Forbidden',
                404: 'Not Found', 405: 'Method Not Allowed', 500: 'Internal Server Error',
                502: 'Bad Gateway', 503: 'Service Unavailable'
            }
            status_text = status_text_map.get(response.status_code, 'OK' if 200 <= response.status_code < 300 else 'Error')
            
            status_icon = "✓" if 200 <= response.status_code < 300 else "✗"
            print(f"[{status_icon}] {case_description}: {method.upper()} {path} -> {response.status_code} {status_text}")
            
            if 200 <= response.status_code < 300:
                success_count += 1
            else:
                fail_count += 1
                
        except requests.exceptions.RequestException as e:
            print(f"[✗] {case_description}: {method.upper()} {path} -> Failed: {e}")
            fail_count += 1
    
    print("=" * 80)
    print(f"\nSummary: {success_count} succeeded, {fail_count} failed, {len(case_ids)} total")
    print()
    
    return 0 if fail_count == 0 else 1


if __name__ == '__main__':
    exit(main())