FlowAnalyzer 0.4.3__tar.gz → 0.4.5__tar.gz

flowanalyzer-0.4.5/FlowAnalyzer/FlowAnalyzer.py

@@ -0,0 +1,303 @@
+import os
+import sqlite3
+import subprocess
+from concurrent.futures import ThreadPoolExecutor
+from typing import Iterable, Optional
+
+from .logging_config import logger
+from .Models import HttpPair, Request, Response
+from .PacketParser import PacketParser
+from .Path import get_default_tshark_path
+
+
+class FlowAnalyzer:
+    """
+    FlowAnalyzer 流量分析器 (智能缓存版)
+    特点：
+    1. Tshark -> Pipe -> ThreadPool -> SQLite
+    2. 智能校验：自动比对 Filter 和文件修改时间，防止缓存错乱
+    3. 存储优化：数据库文件生成在流量包同级目录下
+    """
+
+    def __init__(self, db_path: str):
+        """
+        初始化 FlowAnalyzer
+        :param db_path: 数据库文件路径 (由 get_json_data 返回)
+        """
+        # 路径兼容处理
+        if db_path.endswith(".json"):
+            possible_db = db_path + ".db"
+            if os.path.exists(possible_db):
+                self.db_path = possible_db
+            else:
+                self.db_path = db_path
+        else:
+            self.db_path = db_path
+
+        self.check_db_file()
+
+    def check_db_file(self):
+        """检查数据库文件是否存在"""
+        if not os.path.exists(self.db_path):
+            raise FileNotFoundError(f"未找到数据文件或缓存数据库: {self.db_path}，请先调用 get_json_data 生成。")
+
+    def generate_http_dict_pairs(self) -> Iterable[HttpPair]:
+        """生成HTTP请求和响应信息的字典对 (SQL JOIN 高性能版)"""
+        if not os.path.exists(self.db_path):
+            return
+
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
+            # 开启查询优化
+            cursor.execute("PRAGMA query_only = 1;")
+
+            # === 第一步：配对查询 ===
+            # 利用 SQLite 的 LEFT JOIN 直接匹配请求和响应
+            sql_pair = """
+            SELECT 
+                req.frame_num, req.header, req.file_data, req.full_uri, req.time_epoch,  -- 0-4 (Request)
+                resp.frame_num, resp.header, resp.file_data, resp.time_epoch, resp.request_in -- 5-9 (Response)
+            FROM requests req
+            LEFT JOIN responses resp ON req.frame_num = resp.request_in
+            ORDER BY req.frame_num ASC
+            """
+
+            cursor.execute(sql_pair)
+
+            # 流式遍历结果，内存占用极低
+            for row in cursor:
+                req = Request(frame_num=row[0], header=row[1] or b"", file_data=row[2] or b"", full_uri=row[3] or "", time_epoch=row[4])
+
+                resp = None
+                if row[5] is not None:
+                    resp = Response(frame_num=row[5], header=row[6] or b"", file_data=row[7] or b"", time_epoch=row[8], _request_in=row[9])
+
+                yield HttpPair(request=req, response=resp)
+
+            # === 第二步：孤儿响应查询 ===
+            sql_orphan = """
+            SELECT frame_num, header, file_data, time_epoch, request_in
+            FROM responses
+            WHERE request_in NOT IN (SELECT frame_num FROM requests)
+            """
+            cursor.execute(sql_orphan)
+
+            for row in cursor:
+                resp = Response(frame_num=row[0], header=row[1] or b"", file_data=row[2] or b"", time_epoch=row[3], _request_in=row[4])
+                yield HttpPair(request=None, response=resp)
+
+    # =========================================================================
+    #  静态方法区域：包含校验逻辑和流式处理
+    # =========================================================================
+
+    @staticmethod
+    def get_json_data(file_path: str, display_filter: str, tshark_path: Optional[str] = None) -> str:
+        """
+        获取数据路径 (智能校验版)。
+        """
+        if not os.path.exists(file_path):
+            raise FileNotFoundError("流量包路径不存在：%s" % file_path)
+
+        abs_file_path = os.path.abspath(file_path)
+        pcap_dir = os.path.dirname(abs_file_path)
+        base_name = os.path.splitext(os.path.basename(abs_file_path))[0]
+        db_path = os.path.join(pcap_dir, f"{base_name}.db")
+
+        if FlowAnalyzer._is_cache_valid(db_path, abs_file_path, display_filter):
+            logger.debug(f"缓存校验通过 (Filter匹配且文件未变)，使用缓存: [{db_path}]")
+            return db_path
+        else:
+            logger.debug(f"缓存失效或不存在 (Filter变更或文件更新)，开始重新解析...")
+
+        tshark_path = FlowAnalyzer.get_tshark_path(tshark_path)
+        FlowAnalyzer._stream_tshark_to_db(abs_file_path, display_filter, tshark_path, db_path)
+
+        return db_path
+
+    @staticmethod
+    def get_db_data(file_path: str, display_filter: str, tshark_path: Optional[str] = None) -> str:
+        return FlowAnalyzer.get_json_data(file_path, display_filter, tshark_path)
+
+    @staticmethod
+    def _is_cache_valid(db_path: str, pcap_path: str, current_filter: str) -> bool:
+        if not os.path.exists(db_path) or os.path.getsize(db_path) == 0:
+            return False
+
+        try:
+            current_mtime = os.path.getmtime(pcap_path)
+            current_size = os.path.getsize(pcap_path)
+
+            with sqlite3.connect(db_path) as conn:
+                cursor = conn.cursor()
+                cursor.execute("SELECT filter, pcap_mtime, pcap_size FROM meta_info LIMIT 1")
+                row = cursor.fetchone()
+
+                if not row:
+                    return False
+
+                cached_filter, cached_mtime, cached_size = row
+
+                if cached_filter == current_filter and cached_size == current_size and abs(cached_mtime - current_mtime) < 0.1:
+                    return True
+                else:
+                    logger.debug(f"校验失败: 缓存Filter={cached_filter} vs 当前={current_filter}")
+                    return False
+
+        except sqlite3.OperationalError:
+            return False
+        except Exception as e:
+            logger.warning(f"缓存校验出错: {e}，将重新解析")
+            return False
+
+    @staticmethod
+    def _stream_tshark_to_db(pcap_path: str, display_filter: str, tshark_path: str, db_path: str):
+        """流式解析并存入DB (多线程版)"""
+        if os.path.exists(db_path):
+            os.remove(db_path)
+
+        with sqlite3.connect(db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute("PRAGMA synchronous = OFF")
+            cursor.execute("PRAGMA journal_mode = MEMORY")
+
+            cursor.execute("CREATE TABLE requests (frame_num INTEGER PRIMARY KEY, header BLOB, file_data BLOB, full_uri TEXT, time_epoch REAL)")
+            cursor.execute("CREATE TABLE responses (frame_num INTEGER PRIMARY KEY, header BLOB, file_data BLOB, time_epoch REAL, request_in INTEGER)")
+
+            cursor.execute("""
+                CREATE TABLE meta_info (
+                    id INTEGER PRIMARY KEY, 
+                    filter TEXT, 
+                    pcap_path TEXT, 
+                    pcap_mtime REAL, 
+                    pcap_size INTEGER
+                )
+            """)
+            conn.commit()
+
+        command = [
+            tshark_path,
+            "-r",
+            pcap_path,
+            "-Y",
+            f"({display_filter})",
+            "-T",
+            "fields",
+            "-e",
+            "http.response.code",  # 0
+            "-e",
+            "http.request_in",  # 1
+            "-e",
+            "tcp.reassembled.data",  # 2
+            "-e",
+            "frame.number",  # 3
+            "-e",
+            "tcp.payload",  # 4
+            "-e",
+            "frame.time_epoch",  # 5
+            "-e",
+            "exported_pdu.exported_pdu",  # 6
+            "-e",
+            "http.request.full_uri",  # 7
+            "-e",
+            "tcp.segment.count",  # 8
+            "-E",
+            "header=n",
+            "-E",
+            "separator=/t",
+            "-E",
+            "quote=n",
+            "-E",
+            "occurrence=f",
+        ]
+
+        logger.debug(f"执行 Tshark: {command}")
+        BATCH_SIZE = 2000
+        MAX_PENDING_BATCHES = 20  # 控制内存中待处理的批次数量 (Backpressure)
+
+        # 使用 ThreadPoolExecutor 并行处理数据
+        max_workers = min(32, (os.cpu_count() or 1) + 4)
+
+        process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=os.path.dirname(os.path.abspath(pcap_path)))
+        try:
+            with sqlite3.connect(db_path) as conn:
+                cursor = conn.cursor()
+
+                with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                    current_batch = []
+                    pending_futures = []  # List[Future]
+
+                    def write_results_to_db(results):
+                        """将一批处理好的结果写入数据库"""
+                        if not results:
+                            return
+
+                        db_req_rows = []
+                        db_resp_rows = []
+
+                        for item in results:
+                            if item["type"] == "response":
+                                db_resp_rows.append((item["frame_num"], item["header"], item["file_data"], item["time_epoch"], item["request_in"]))
+                            else:
+                                db_req_rows.append((item["frame_num"], item["header"], item["file_data"], item["full_uri"], item["time_epoch"]))
+
+                        if db_req_rows:
+                            cursor.executemany("INSERT OR REPLACE INTO requests VALUES (?,?,?,?,?)", db_req_rows)
+                        if db_resp_rows:
+                            cursor.executemany("INSERT OR REPLACE INTO responses VALUES (?,?,?,?,?)", db_resp_rows)
+
+                    def submit_batch():
+                        """提交当前批次到线程池"""
+                        if not current_batch:
+                            return
+
+                        # Copy batch data for the thread (list slicing is fast)
+                        batch_data = current_batch[:]
+                        future = executor.submit(PacketParser.process_batch, batch_data)
+                        pending_futures.append(future)
+                        current_batch.clear()
+
+                    # --- Main Pipeline Loop ---
+                    if process.stdout:
+                        for line in process.stdout:
+                            current_batch.append(line)
+
+                            if len(current_batch) >= BATCH_SIZE:
+                                submit_batch()
+
+                                # Backpressure: 如果积压的任务太多，主线程暂停读取，先处理掉最早的一个
+                                # 这样既保证了 Pipeline 流动，又防止内存爆掉
+                                if len(pending_futures) >= MAX_PENDING_BATCHES:
+                                    oldest_future = pending_futures.pop(0)
+                                    write_results_to_db(oldest_future.result())
+
+                    # --- Drain Pipeline ---
+                    # 提交剩余数据
+                    submit_batch()
+
+                    # 等待所有剩余任务完成
+                    for future in pending_futures:
+                        write_results_to_db(future.result())
+
+                # 创建索引和元数据
+                cursor.execute("CREATE INDEX idx_resp_req_in ON responses(request_in)")
+                pcap_mtime = os.path.getmtime(pcap_path)
+                pcap_size = os.path.getsize(pcap_path)
+                cursor.execute("INSERT INTO meta_info (filter, pcap_path, pcap_mtime, pcap_size) VALUES (?, ?, ?, ?)", (display_filter, pcap_path, pcap_mtime, pcap_size))
+                conn.commit()
+
+        except Exception as e:
+            logger.error(f"解析错误: {e}")
+            if process.poll() is None:
+                process.terminate()
+        finally:
+            if process.poll() is None:
+                process.terminate()
+
+    @staticmethod
+    def get_tshark_path(tshark_path: Optional[str]) -> str:
+        default_tshark_path = get_default_tshark_path()
+        use_path = tshark_path if tshark_path and os.path.exists(tshark_path) else default_tshark_path
+        if not use_path or not os.path.exists(use_path):
+            logger.critical("未找到 Tshark，请检查路径配置")
+            exit(-1)
+        return use_path

flowanalyzer-0.4.5/FlowAnalyzer/Models.py

@@ -0,0 +1,27 @@
+from dataclasses import dataclass
+from typing import NamedTuple, Optional
+
+
+@dataclass
+class Request:
+    __slots__ = ("frame_num", "header", "file_data", "full_uri", "time_epoch")
+    frame_num: int
+    header: bytes
+    file_data: bytes
+    full_uri: str
+    time_epoch: float
+
+
+@dataclass
+class Response:
+    __slots__ = ("frame_num", "header", "file_data", "time_epoch", "_request_in")
+    frame_num: int
+    header: bytes
+    file_data: bytes
+    time_epoch: float
+    _request_in: Optional[int]
+
+
+class HttpPair(NamedTuple):
+    request: Optional[Request]
+    response: Optional[Response]

flowanalyzer-0.4.5/FlowAnalyzer/PacketParser.py

@@ -0,0 +1,183 @@
+import binascii
+import contextlib
+import gzip
+from typing import List, Optional, Tuple
+from urllib import parse
+
+from .logging_config import logger
+
+
+class PacketParser:
+    @staticmethod
+    def parse_packet_data(row: list) -> Tuple[int, int, float, str, bytes]:
+        """
+        解析 Tshark 输出的一行数据
+        row definition (all bytes):
+        0: http.response.code
+        1: http.request_in
+        2: tcp.reassembled.data
+        3: frame.number
+        4: tcp.payload
+        5: frame.time_epoch
+        6: exported_pdu.exported_pdu
+        7: http.request.full_uri 
+        8: tcp.segment.count
+        """
+        frame_num = int(row[3])
+        request_in = int(row[1]) if row[1] else frame_num
+        # Decode only URI to string
+        full_uri = parse.unquote(row[7].decode("utf-8", errors="replace")) if row[7] else ""
+        time_epoch = float(row[5])
+
+        # Logic for Raw Packet (Header Source)
+        # Previous index 9 is now 8 since we removed http.file_data
+        is_reassembled = len(row) > 8 and row[8]
+
+        if is_reassembled and row[2]:
+            full_request = row[2]
+        elif row[4]:
+            full_request = row[4]
+        else:
+            # Fallback (e.g. Exported PDU)
+            full_request = row[2] if row[2] else (row[6] if row[6] else b"")
+
+        return frame_num, request_in, time_epoch, full_uri, full_request
+
+    @staticmethod
+    def split_http_headers(file_data: bytes) -> Tuple[bytes, bytes]:
+        headerEnd = file_data.find(b"\r\n\r\n")
+        if headerEnd != -1:
+            return file_data[: headerEnd + 4], file_data[headerEnd + 4 :]
+        elif file_data.find(b"\n\n") != -1:
+            headerEnd = file_data.index(b"\n\n") + 2
+            return file_data[:headerEnd], file_data[headerEnd:]
+        return b"", file_data
+
+    @staticmethod
+    def dechunk_http_response(file_data: bytes) -> bytes:
+        """解码分块TCP数据"""
+        if not file_data:
+            return b""
+
+        chunks = []
+        cursor = 0
+        total_len = len(file_data)
+
+        while cursor < total_len:
+            newline_idx = file_data.find(b"\n", cursor)
+            if newline_idx == -1:
+                # If no newline found, maybe it's just remaining data (though strictly should end with 0 chunk)
+                # But for robustness we might perform a "best effort" or just stop.
+                # raising ValueError("Not chunked data") might be too aggressive if we are just "trying" to dechunk
+                # Let's assume non-chunked if strict format not found
+                raise ValueError("Not chunked data")
+
+            size_line = file_data[cursor:newline_idx].strip()
+            if not size_line:
+                cursor = newline_idx + 1
+                continue
+
+            try:
+                chunk_size = int(size_line, 16)
+            except ValueError:
+                raise ValueError("Invalid chunk size")
+
+            if chunk_size == 0:
+                break
+
+            data_start = newline_idx + 1
+            data_end = data_start + chunk_size
+
+            # Robustness check
+            if data_start > total_len:
+                break
+
+            if data_end > total_len:
+                chunks.append(file_data[data_start:])
+                break
+
+            chunks.append(file_data[data_start:data_end])
+
+            cursor = data_end
+            # Skip CRLF after chunk data
+            while cursor < total_len and file_data[cursor] in (13, 10):
+                cursor += 1
+
+        return b"".join(chunks)
+
+    @staticmethod
+    def extract_http_file_data(full_request: bytes) -> Tuple[bytes, bytes]:
+        """
+        提取HTTP请求或响应中的文件数据 (混合模式 - 二进制优化版)
+        """
+        header = b""
+        file_data = b""
+
+        if not full_request:
+            return b"", b""
+        try:
+            raw_bytes = binascii.unhexlify(full_request)
+            header, body_part = PacketParser.split_http_headers(raw_bytes)
+
+            with contextlib.suppress(Exception):
+                body_part = PacketParser.dechunk_http_response(body_part)
+
+            with contextlib.suppress(Exception):
+                if body_part.startswith(b"\x1f\x8b"):
+                    body_part = gzip.decompress(body_part)
+
+            file_data = body_part
+            return header, file_data
+
+        except binascii.Error:
+            logger.error("Hex转换失败")
+            return b"", b""
+        except Exception as e:
+            logger.error(f"解析HTTP数据未知错误: {e}")
+            return b"", b""
+
+    @staticmethod
+    def process_row(line: bytes) -> Optional[dict]:
+        """
+        处理单行数据，返回结构化结果供主线程写入
+        """
+        line = line.rstrip(b"\r\n")
+        if not line:
+            return None
+
+        row = line.split(b"\t")
+        try:
+            frame_num, request_in, time_epoch, full_uri, full_request = PacketParser.parse_packet_data(row)
+
+            if not full_request:
+                return None
+
+            header, file_data = PacketParser.extract_http_file_data(full_request)
+
+            # row[0] is http.response.code (bytes)
+            is_response = bool(row[0])
+
+            return {
+                "type": "response" if is_response else "request",
+                "frame_num": frame_num,
+                "header": header,
+                "file_data": file_data,
+                "time_epoch": time_epoch,
+                "request_in": request_in,  # Only useful for Response
+                "full_uri": full_uri,  # Only useful for Request
+            }
+
+        except Exception:
+            return None
+
+    @staticmethod
+    def process_batch(lines: List[bytes]) -> List[dict]:
+        """
+        批量处理行数据，减少函数调用开销
+        """
+        results = []
+        for line in lines:
+            res = PacketParser.process_row(line)
+            if res:
+                results.append(res)
+        return results

flowanalyzer-0.4.5/FlowAnalyzer/PcapSplitter.py

@@ -0,0 +1,128 @@
+import os
+import time
+from collections import defaultdict
+from typing import List, Tuple
+
+import dpkt
+
+
+class PcapSplitter:
+    """
+    Encapsulates logic to split a PCAP/PCAPNG file into multiple smaller PCAP files
+    based on TCP flows, dynamically balanced for parallel processing.
+    """
+
+    def __init__(self, pcap_file: str, output_dir: str):
+        self.pcap_file = pcap_file
+        self.output_dir = output_dir
+
+    def get_stream_key(self, tcp, ip) -> Tuple:
+        """Generate a 5-tuple key for the flow."""
+        src = ip.src
+        dst = ip.dst
+        sport = tcp.sport
+        dport = tcp.dport
+        # Canonicalize bidirectional flows to the same key
+        key1 = (src, dst, sport, dport)
+        key2 = (dst, src, dport, sport)
+        return key1 if key1 < key2 else key2
+
+    def split(self, threshold_mb: int = 10, default_chunks: int = 3) -> List[str]:
+        """
+        Split the pcap file into balanced chunks based on stream volume (bytes).
+        Uses a Greedy Partition Algorithm (Longest Processing Time first).
+
+        Args:
+            threshold_mb: File size threshold in MB. If smaller, do not split.
+            default_chunks: Number of chunks to split into if threshold is exceeded.
+
+        Returns:
+            List of generated file paths (or original file if not split).
+        """
+        if not os.path.exists(self.pcap_file):
+            raise FileNotFoundError(f"PCAP file not found: {self.pcap_file}")
+
+        file_size_mb = os.path.getsize(self.pcap_file) / (1024 * 1024)
+        if file_size_mb < threshold_mb:
+            print(f"File size {file_size_mb:.2f}MB < {threshold_mb}MB. Skipping split.")
+            return [self.pcap_file]
+
+        os.makedirs(self.output_dir, exist_ok=True)
+
+        start_time = time.time()
+        # Dictionary to store packets: stream_key -> list of (ts, buf)
+        streams = defaultdict(list)
+        # Dictionary to store total size: stream_key -> total_bytes
+        stream_sizes = defaultdict(int)
+
+        # 1. Read and Group Packets
+        print(f"Reading {self.pcap_file}...")
+        with open(self.pcap_file, "rb") as f:
+            if self.pcap_file.lower().endswith(".pcapng"):
+                reader = dpkt.pcapng.Reader(f)
+            else:
+                reader = dpkt.pcap.Reader(f)
+
+            for ts, buf in reader:
+                try:
+                    eth = dpkt.ethernet.Ethernet(buf)
+                    if not isinstance(eth.data, dpkt.ip.IP):
+                        continue
+                    ip = eth.data
+                    if not isinstance(ip.data, dpkt.tcp.TCP):
+                        continue
+                    tcp = ip.data
+
+                    key = self.get_stream_key(tcp, ip)
+                    streams[key].append((ts, buf))
+                    stream_sizes[key] += len(buf)
+                except Exception:
+                    continue
+
+        total_streams = len(streams)
+        print(f"Found {total_streams} TCP streams.")
+
+        if total_streams == 0:
+            print("No TCP streams found to split.")
+            return []
+
+        # 2. Assign Streams to Buckets (Greedy LPT Algorithm)
+        num_chunks = min(default_chunks, total_streams)
+
+        # Sort streams by size (descending)
+        sorted_streams = sorted(stream_sizes.items(), key=lambda item: item[1], reverse=True)
+
+        # Buckets: list of (current_size, batch_index, list_of_keys)
+        # We perform standard list sort to find min bucket, sufficient for small N
+        buckets = [[0, i, []] for i in range(num_chunks)]
+
+        for key, size in sorted_streams:
+            # Find bucket with smallest current size
+            buckets.sort(key=lambda x: x[0])
+            smallest_bucket = buckets[0]
+
+            # Add stream to this bucket
+            smallest_bucket[0] += size
+            smallest_bucket[2].append(key)
+
+        print(f"Splitting into {num_chunks} files with volume balancing...")
+        generated_files = []
+
+        # 3. Write Batches
+        # Sort buckets by index ensures file naming order 0, 1, 2...
+        buckets.sort(key=lambda x: x[1])
+
+        for size, i, batch_keys in buckets:
+            out_file_path = os.path.join(self.output_dir, f"batch_{i}.pcap")
+            generated_files.append(out_file_path)
+
+            with open(out_file_path, "wb") as f:
+                writer = dpkt.pcap.Writer(f)
+                for key in batch_keys:
+                    for ts, buf in streams[key]:
+                        writer.writepkt(buf, ts)
+
+            print(f"  - Created {os.path.basename(out_file_path)}: {len(batch_keys)} streams ({size/1024/1024:.2f} MB)")
+
+        print(f"Split completed in {time.time() - start_time:.2f}s")
+        return generated_files

{flowanalyzer-0.4.3 → flowanalyzer-0.4.5}/FlowAnalyzer/logging_config.py

@@ -15,8 +15,9 @@ def configure_logger(logger_name, level=logging.DEBUG) -> logging.Logger:
     console_handler.setFormatter(formatter)
     return logger
 
+
 logger = configure_logger("FlowAnalyzer", logging.INFO)
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     logger = configure_logger("FlowAnalyzer")
     logger.info("This is a test!")