import copy
import re

from scapy.layers.inet import TCP
from scapy.packet import Raw

from xbase_util.xbase_constant import plain_content_type_columns, packetKeyname, src_dst_header, statisticHeader, \
    features_key, plain_body_columns
from xbase_util.xbase_util import firstOrZero


def content_type_is_plain(packet):
    """
    从单个包（包括header和body）中获取content-type并判断是否为可见类型
    :param packet:
    :return:
    """
    if ":" not in packet:
        return False
    for item in packet.replace("-", "_").replace(" ", "").lower().splitlines():
        if "content_type" in item:
            if ":" not in item:
                continue
            content_type = item.split(":")[1].replace("\r", "").strip()
            return content_type in plain_content_type_columns
    return False


def filter_visible_chars(data):
    """
    过滤不可见字符，仅保留可打印的ASCII字符
    :param data:
    :return:
    """
    return ''.join(chr(b) for b in data if 32 <= b <= 126 or b in (9, 10, 13))


def get_all_columns(
        contains_packet_column=False,
        contains_src_dst_column=False,
        contains_statistic_column=False,
        contains_features_column=False,
        contains_plain_body_column=False,
        contains_pcap_flow_text=False
):
    result_columns = []
    if contains_packet_column:
        result_columns += packetKeyname
    if contains_src_dst_column:
        result_columns += src_dst_header
    if contains_statistic_column:
        result_columns += statisticHeader
    if contains_features_column:
        result_columns += features_key
    if contains_plain_body_column:
        result_columns += plain_body_columns
    if contains_pcap_flow_text:
        result_columns.append(contains_pcap_flow_text)
    return result_columns


req_pattern = re.compile(r"(GET|POST|HEAD|PUT|DELETE|OPTIONS|PATCH) \/[^\s]* HTTP\/\d\.\d[\s\S]*?\r\n\r\n",
                         re.DOTALL)
res_pattern = re.compile(r"HTTP/\d\.\d \d{3}.*", re.DOTALL)
req_body_pattern = re.compile(
    r"(GET|POST|HEAD|PUT|DELETE|OPTIONS|PATCH) \/[^\s]* HTTP\/\d\.\d[\s\S]*?(?=HTTP/\d\.\d)", re.DOTALL)

http_version = re.compile(r"HTTP\/(\d\.\d)")
http_req_method = re.compile(r"(GET|POST|HEAD|PUT|DELETE|OPTIONS|PATCH) \/[^\s]* HTTP\/\d\.\d")
http_req_path = re.compile(r"(?:GET|POST|HEAD|PUT|DELETE|OPTIONS|PATCH)\s+(\/[^\s]*)\s+HTTP\/\d\.\d")
res_status_code_pattern = re.compile(r"HTTP\/\d\.\d\s+(\d{3})\s+.*")


# def get_all_packets_by_reg(packets):
#     http_Req_Raw = {}
#     http_methods = ("POST /", "PUT /", "OPTIONS /", "DELETE /", "GET /")
#     for packet in packets:
#         if TCP in packet and Raw in packet:
#             data = packet[Raw].load
#             ack = packet[TCP].ack
#             next_ack = packet[TCP].seq + len(data)
#             if ack not in http_Req_Raw:
#                 http_Req_Raw[ack] = {"time": [float(packet.time)], "data": data, "next_ack": next_ack}
#             else:
#                 http_Req_Raw[ack]["time"].append(float(packet.time))
#                 http_Req_Raw[ack]["data"] += data
#                 http_Req_Raw[ack]["next_ack"] = next_ack
#     packet_list = [
#         {
#             'req_data': item['data'],
#             'res_data': http_Req_Raw[item['next_ack']]['data'],
#             'req_text': filter_visible_chars(item['data']),
#             'res_text': filter_visible_chars(http_Req_Raw[item['next_ack']]['data']),
#             'req_time': item['time'],
#             'res_time': http_Req_Raw[item['next_ack']]['time']
#         }
#         for ack, item in http_Req_Raw.items()
#         if item['next_ack'] in http_Req_Raw  # 检查 next_ack 是否存在
#            and any(method in filter_visible_chars(item['data']) for method in http_methods)
#     ]
#     return packet_list


def get_header_value(header_set, value):
    result = [item for item in header_set if value in item]
    if len(result) != 0:
        return result[0].replace(f"{value}:", "").strip()
    else:
        return ""


def get_detail_by_package(publicField, req_header, req_body, res_header, res_body):
    """
    通过pcap的数量分离session并完善相关字段
    :param publicField: 原始的session单条数据
    :param req_header:请求头
    :param req_body:请求体
    :param res_header:响应头
    :param res_body:响应体
    :return: 完整的单条数据
    """
    res_field = copy.deepcopy(publicField)
    res_field["initRTT"] = firstOrZero(res_field.get("initRTT", 0))
    res_field["length"] = firstOrZero(res_field.get("length", 0))

    http_version_res = http_version.findall(req_header)
    res_field['http.clientVersion'] = http_version_res[0] if len(http_version_res) > 0 else ""
    http_method = http_req_method.findall(req_header)
    http_path = http_req_path.findall(req_header)
    res_field['http.clientVersion'] = http_version_res[0] if len(http_version_res) > 0 else ""
    res_field['http.method'] = http_method[0] if len(http_method) > 0 else ""
    res_field['http.path'] = http_path[0] if len(http_path) > 0 else ""
    request_lines = req_header.splitlines()
    res_field['http.request-referer'] = get_header_value(header_set=request_lines, value="Referer")
    res_field['http.request-content-type'] = get_header_value(header_set=request_lines,
                                                              value="Content-Type")
    res_field['http.hostTokens'] = get_header_value(header_set=request_lines, value="Host")

    res_field['plain_body_src'] = ""
    res_field['plain_body_dst'] = ""
    if content_type_is_plain(req_header):
        res_field['plain_body_src'] = filter_visible_chars(f'{req_body}')
    if content_type_is_plain(res_header):
        res_field['plain_body_dst'] = filter_visible_chars(f'{res_body}')

    http_server_version_res = http_version.findall(res_header)
    res_field['http.serverVersion'] = http_server_version_res[0] if len(http_server_version_res) > 0 else ""

    status_code = res_status_code_pattern.findall(res_header)
    res_field['http.statuscode'] = status_code[0] if len(status_code) > 0 else ""
    response_lines = res_header.splitlines()
    res_field['http.response-server'] = get_header_value(header_set=response_lines, value="Server")
    res_field['http.response-content-type'] = get_header_value(header_set=response_lines,
                                                               value="Content-Type")
    for response in list(set(response_lines + request_lines)):
        key_value = response.replace("\r", "").split(":")
        if len(key_value) == 2:
            key = key_value[0].replace(" ", "").replace("-", "_").lower()
            value = key_value[1].replace(" ", "")
            if f"src_{key}" in src_dst_header:
                res_field[f"src_{key}"] = value
            if f"dst_{key}" in src_dst_header:
                res_field[f"dst_{key}"] = value
    return res_field
