# -*- coding: utf-8 -*-
"""
Усовершенствованный анализатор логов для Privoxy с поддержкой IPv6,
отслеживанием соединений, цветовым кодированием и гибкой фильтрацией.

Этот скрипт может работать в двух режимах:
1.  --live: отслеживание файла лога в реальном времени (аналог 'tail -f').
2.  --full: разбор указанного количества последних строк из файла и выход.

Основные возможности:
- Корреляция логов 'Request' и 'Connected to' по PID для отображения полной цепочки.
- Поддержка различных форматов логов, включая Crunch-ошибки и Apache-совместимый формат.
- Корректная обработка и отображение IPv4 и IPv6 адресов.
- Цветовое кодирование вывода для лучшей читаемости.
- Фильтрация вывода по доменам из файла.
- Режимы отладки, отображения только ошибок или всех записей.
"""

import argparse
import time
import re
import os
from datetime import datetime
from collections import deque, OrderedDict
from colorama import init, Fore, Back, Style

# Инициализация colorama для поддержки цветов в терминале
init()

# ====================== НАСТРОЙКИ ЦВЕТОВ ======================
class LogColors:
    """Класс для хранения кодов цветов ANSI."""
    TIMESTAMP = Fore.WHITE
    CLIENT_IP = Fore.LIGHTYELLOW_EX
    CONNECT = Fore.LIGHTYELLOW_EX
    DIRECT = Fore.LIGHTGREEN_EX
    DEBUG = Fore.CYAN
    CRUNCH_NORMAL = Fore.LIGHTWHITE_EX
    CRUNCH_HIGHLIGHT = f"{Fore.BLACK}{Back.WHITE}{Style.BRIGHT}"
    HTTP_200 = Style.DIM
    HTTP_300 = Fore.LIGHTBLUE_EX
    HTTP_400 = Fore.YELLOW
    HTTP_500 = Fore.LIGHTRED_EX
    ERROR = Fore.LIGHTRED_EX
    
    @classmethod
    def disable_colors(cls):
        """Отключает цвета, заменяя коды на пустые строки."""
        for attr in dir(cls):
            if not attr.startswith('__') and not callable(getattr(cls, attr)):
                setattr(cls, attr, '')

# Глобальный флаг для управления цветами
USE_COLORS = True

def enable_colors(enable):
    """Включает или отключает цветной вывод."""
    global USE_COLORS
    USE_COLORS = enable
    if not enable:
        LogColors.disable_colors()
# ==============================================================

class DomainFilter:
    """
    Фильтрует логи на основе списка доменов.
    Если хост или любой из его родительских доменов находится в списке,
    он считается отфильтрованным.
    """
    def __init__(self, filter_file=None):
        self.domains = set()
        if filter_file and os.path.exists(filter_file):
            with open(filter_file, 'r', encoding='utf-8') as f:
                for line in f:
                    domain = line.strip()
                    if domain and not domain.startswith('#'):
                        # Отбрасываем порт, если он есть
                        domain = domain.split(':')[0]
                        self.domains.add(domain.lower())
    
    def is_filtered(self, host):
        """Проверяет, должен ли хост быть отфильтрован."""
        if not host or not self.domains:
            return False
        
        host_without_port = host.split(':')[0].lower()
        
        if host_without_port in self.domains:
            return True
        
        # Проверяем родительские домены (например, 'google.com' для 'www.google.com')
        parts = host_without_port.split('.')
        for i in range(1, len(parts)):
            sub_domain = '.'.join(parts[i:])
            if sub_domain in self.domains:
                return True
        
        return False

# ====================== РЕГУЛЯРНЫЕ ВЫРАЖЕНИЯ ======================

# Новое: для парсинга IP-адреса клиента
accepted_conn_re = re.compile(
    r'^(?P<dt>\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2})\.\d+\s+(?P<pid>\S+)\s+'
    r'Connect: Accepted connection from (?P<client_ip>\S+) on socket \d+$'
)

# Crunch-ошибки (например, "forwarding failed")
crunch_re = re.compile(
    r'(?P<dt>\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2})\.\d+\s+\S+\s+Crunch:\s+'
    r'(?P<msg>[^:]+):\s*'
    r'(?P<host>.+)',
    re.IGNORECASE
)

# Apache-совместимый формат лога
apache_re = re.compile(
    r'^(\S+)\s+- - \[([^\]]+)\] "(\w+) ([^"]+)" (\d{3}) (\d+|-)'
)

# --- Регулярные выражения для отслеживания соединений ---

# Шаблон 'Request:': начало отслеживания по PID
request_re = re.compile(
    r'^(?P<dt>\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3})\s+(?P<pid>\S+)\s+'
    r'Request: (?P<host>\[[^\]]+\]|[^/:\s]+)(?::(?P<host_port>\d+))?(?:/|$)'
)

# Шаблон 'Connected to': используется для сопоставления с 'Request'
connected_to_re = re.compile(
    r'^(?P<dt>\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3})\s+(?P<pid>\S+)\s+'
    r'Connect: Connected to (?P<via>[^\s\[]+)(?:\[(?P<via_addr>[^\]]+)\])?(?::(?P<via_port>\d+))?\.$'
)

# Шаблоны 'Connect: via': обрабатываются немедленно, без сопоставления
VIA_CONNECT_PARSERS = {
    'brackets_with_port': re.compile(
        r'^(?P<dt>\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3})\s+(?P<pid>\S+)\s+'
        r'Connect: via \[(?P<via>[^]]+)\]:?(?P<via_port>\d+)?\]\s+'
        r'to:\s+(?P<host>\[[^\]]+\]|[^:\s]+)(?::(?P<host_port>\d+))?$'
    ),
    'no_brackets_with_port': re.compile(
        r'^(?P<dt>\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3})\s+(?P<pid>\S+)\s+'
        r'Connect: via\s+(?P<via>[^:\s]+):(?P<via_port>\d+)\s+'
        r'to:\s+(?P<host>\[[^\]]+\]|[^:\s]+)(?::(?P<host_port>\d+))?$'
    ),
    'no_brackets_without_port': re.compile(
        r'^(?P<dt>\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3})\s+(?P<pid>\S+)\s+'
        r'Connect: via\s+(?P<via>[^:\s]+)\s+'
        r'to:\s+(?P<host>\[[^\]]+\]|[^:\s]+)(?::(?P<host_port>\d+))?$'
    ),
}
# =================================================================

class ConnectionTracker:
    """
    Отслеживает и сопоставляет разрозненные лог-записи о соединениях по PID.
    """
    def __init__(self, max_size=1000, ttl=3600):
        self.requests = OrderedDict()
        self.connections_proxy = OrderedDict()
        self.handled_by_via = OrderedDict()
        self.client_ips = OrderedDict()
        self.max_size = max_size
        self.ttl = ttl  # Время жизни записи в секундах

    def _cleanup(self, storage):
        """Удаляет старые записи из указанного хранилища."""
        current_time = time.time()
        expired_keys = [
            pid for pid, record in storage.items()
            if current_time - record.get('_timestamp', 0) > self.ttl
        ]
        for pid in expired_keys:
            storage.pop(pid, None)
        while len(storage) > self.max_size:
            storage.popitem(last=False)

    def _get_timestamp(self, dt_str):
        """Преобразует строку времени из лога в Unix timestamp."""
        try:
            dt_obj = datetime.strptime(dt_str.split('.')[0], "%Y-%m-%d %H:%M:%S")
            return dt_obj.timestamp()
        except (ValueError, IndexError):
            return time.time()

    def add_request(self, pid, host, host_port, dt):
        self._cleanup(self.requests)
        self.requests[pid] = {
            'host': host, 'host_port': host_port, 'dt': dt,
            '_timestamp': self._get_timestamp(dt)
        }

    def add_connection_proxy(self, pid, via, via_addr, via_port, dt):
        self._cleanup(self.connections_proxy)
        self.connections_proxy[pid] = {
            'via': via, 'via_addr': via_addr, 'via_port': via_port, 'dt': dt, 'shown': False,
            '_timestamp': self._get_timestamp(dt)
        }
    
    def add_client_ip(self, pid, client_ip):
        """Сохраняет IP-адрес клиента для данного PID."""
        self._cleanup(self.client_ips)
        self.client_ips[pid] = {
            'ip': client_ip,
            '_timestamp': time.time()
        }

    def get_client_ip(self, pid):
        """Получает IP-адрес клиента для данного PID."""
        record = self.client_ips.get(pid)
        return record['ip'] if record else '???.???.???.???'

    def add_via_handled_pid(self, pid):
        """Отмечает, что для PID была выведена строка 'via'."""
        self._cleanup(self.handled_by_via)
        self.handled_by_via[pid] = {'_timestamp': time.time()}

    def is_via_handled(self, pid):
        return pid in self.handled_by_via

    def get_connection_info(self, pid):
        """
        Пытается найти пару "запрос-соединение".
        Подавляет вывод для прокси-соединений, которые уже были
        обработаны парсером 'Connect: via', чтобы избежать дублирования.
        """
        if self.is_via_handled(pid):
            return None

        request = self.requests.get(pid)
        connection_proxy = self.connections_proxy.get(pid)

        if request and connection_proxy and not connection_proxy['shown']:
            connection_proxy['shown'] = True
            is_direct = (request['host'].lower() == connection_proxy['via'].lower())
            return {
                'dt': connection_proxy['dt'], 'pid': pid,
                'via': connection_proxy['via'], 'via_addr': connection_proxy['via_addr'],
                'via_port': connection_proxy['via_port'], 'host': request['host'],
                'host_port': request['host_port'], 'is_direct': is_direct
            }
        return None

# Глобальный экземпляр трекера
connection_tracker = ConnectionTracker()

def _format_host_with_port(host, port):
    """Корректно форматирует хост и порт, добавляя скобки для IPv6."""
    if not port:
        return host
    if ':' in host and not host.startswith('['):
        return f"[{host}]:{port}"
    return f"{host}:{port}"

def convert_apache_time(apache_timestamp):
    """Преобразует время из формата Apache в 'YYYY-MM-DD HH:MM:SS'."""
    try:
        dt = datetime.strptime(apache_timestamp.split()[0], "%d/%b/%Y:%H:%M:%S")
        return dt.strftime("%Y-%m-%d %H:%M:%S")
    except ValueError:
        return apache_timestamp.split()[0]

def get_last_lines(file_path, num_lines=1000):
    try:
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            return list(deque(f, maxlen=num_lines))
    except (IOError, FileNotFoundError) as e:
        print(f"{LogColors.ERROR}Ошибка чтения файла: {e}{Style.RESET_ALL if USE_COLORS else ''}")
        return []

def parse_connect_via_line(line, debug=False):
    """Парсит строки 'Connect: via ...'."""
    clean_line = line.strip()
    for parser_name, parser in VIA_CONNECT_PARSERS.items():
        match = parser.match(clean_line)
        if match:
            if debug:
                print(f"{LogColors.DEBUG}MATCHED via-parser: {parser_name}{Style.RESET_ALL if USE_COLORS else ''}")
            groups = match.groupdict()
            return {
                'dt': groups['dt'], 'pid': groups['pid'], 'via': groups['via'].strip('[]'),
                'via_port': groups.get('via_port'), 'host': groups['host'].strip('[]'),
                'host_port': groups.get('host_port')
            }
    return None

def parse_line(line, show_all=False, error_only=False, debug=False, domain_filter=None):
    """
    Главная функция-парсер. Определяет тип строки и форматирует ее для вывода.
    """
    if debug:
        print(f"\n{LogColors.DEBUG}=== DEBUG RAW LINE: {repr(line.strip())} ==={Style.RESET_ALL if USE_COLORS else ''}")

    line = line.strip()
    if not line:
        return None

    # --- 0. Обработка 'Accepted connection' для получения IP клиента ---
    if match := accepted_conn_re.match(line):
        groups = match.groupdict()
        connection_tracker.add_client_ip(groups['pid'], groups['client_ip'])
        if debug:
            print(f"{LogColors.DEBUG}Registered client IP: PID={groups['pid']}, IP={groups['client_ip']}{Style.RESET_ALL if USE_COLORS else ''}")
        return None

    # --- 1. Обработка 'Request:' ---
    if match := request_re.match(line):
        groups = match.groupdict()
        host = groups['host'].strip('[]')
        if domain_filter and domain_filter.is_filtered(host):
            return None
        connection_tracker.add_request(groups['pid'], host, groups.get('host_port'), groups['dt'])
        if debug:
            print(f"{LogColors.DEBUG}Registered request: PID={groups['pid']}, host={host}{Style.RESET_ALL if USE_COLORS else ''}")
        return None

    # --- 2. Обработка 'Connected to:' ---
    if match := connected_to_re.match(line):
        if error_only: return None
        groups = match.groupdict()
        connection_tracker.add_connection_proxy(
            groups['pid'], groups['via'], groups.get('via_addr'), groups.get('via_port'), groups['dt']
        )
        if info := connection_tracker.get_connection_info(groups['pid']):
            if domain_filter and domain_filter.is_filtered(info['host']):
                return None
            
            pid = info['pid']
            dt = info['dt'].split('.')[0]
            host_full = _format_host_with_port(info['host'], info['host_port'])
            client_ip = connection_tracker.get_client_ip(pid)
            client_ip_colored = f"{LogColors.CLIENT_IP}{client_ip}{Style.RESET_ALL if USE_COLORS else ''}"
            
            if info['is_direct']:
                return f"{LogColors.TIMESTAMP}{dt}{Style.RESET_ALL if USE_COLORS else ''} {client_ip_colored} {LogColors.DIRECT}DIRECT {host_full} (PID: {pid}){Style.RESET_ALL if USE_COLORS else ''}"
            else:
                via_full = _format_host_with_port(info['via_addr'] or info['via'], info['via_port'])
                return f"{LogColors.TIMESTAMP}{dt}{Style.RESET_ALL if USE_COLORS else ''} {client_ip_colored} {LogColors.CONNECT}CONNECT {host_full} proxy [{via_full}] (PID: {pid}){Style.RESET_ALL if USE_COLORS else ''}"
        return None

    # --- 3. Обработка 'Crunch:' (ошибки) ---
    if match := crunch_re.match(line):
        groups = match.groupdict()
        host = groups["host"].strip()
        if domain_filter and domain_filter.is_filtered(host): return None
        dt, msg = groups["dt"], groups["msg"].strip()
        color = LogColors.CRUNCH_HIGHLIGHT if show_all else LogColors.CRUNCH_NORMAL
        return f"{LogColors.TIMESTAMP}{dt}{Style.RESET_ALL if USE_COLORS else ''} {color}Crunch: {msg}: {host}{Style.RESET_ALL if USE_COLORS else ''}"

    # --- 4. Обработка Apache-формата ---
    if match := apache_re.match(line):
        try:
            ip, timestamp, method, url, code_str, size = match.groups()
            code = int(code_str)
            if error_only and code < 400: return None
            
            host = url.split('/')[0].strip('[]') if '/' in url else url.strip('[]')
            if domain_filter and domain_filter.is_filtered(host): return None
            
            color = ''
            if code >= 500: color = LogColors.HTTP_500
            elif code >= 400: color = LogColors.HTTP_400
            elif code >= 300: color = LogColors.HTTP_300
            elif show_all and code in [200, 206]: color = LogColors.HTTP_200
            if not color and not show_all and not error_only: return None

            formatted_time = convert_apache_time(timestamp)
            result = f"{LogColors.TIMESTAMP}{formatted_time}{Style.RESET_ALL if USE_COLORS else ''} {ip} {method} {url} {code} {size}"
            return f"{color}{result}{Style.RESET_ALL if USE_COLORS else ''}"
        except (ValueError, IndexError):
            return None

    # --- 5. Обработка 'Connect: via' (для HTTP-прокси) ---
    if "Connect: via" in line and not error_only:
        if match := parse_connect_via_line(line, debug):
            if domain_filter and domain_filter.is_filtered(match['host']): return None
            
            pid = match['pid']
            connection_tracker.add_via_handled_pid(pid)
            
            client_ip = connection_tracker.get_client_ip(pid)
            client_ip_colored = f"{LogColors.CLIENT_IP}{client_ip}{Style.RESET_ALL if USE_COLORS else ''}"
            dt = match['dt'].split('.')[0]
            via_full = _format_host_with_port(match['via'], match['via_port'])
            host_full = _format_host_with_port(match['host'], match['host_port'])

            return f"{LogColors.TIMESTAMP}{dt}{Style.RESET_ALL if USE_COLORS else ''} {client_ip_colored} {LogColors.CONNECT}CONNECT {host_full} via [{via_full}] (PID: {pid}){Style.RESET_ALL if USE_COLORS else ''}"

    if debug:
        print(f"{LogColors.DEBUG}Строка не соответствует ни одному шаблону.{Style.RESET_ALL if USE_COLORS else ''}")
    return None

def follow(path, show_all, error_only, debug, num_lines, domain_filter):
    try:
        for line in get_last_lines(path, num_lines):
            if parsed := parse_line(line, show_all, error_only, debug, domain_filter):
                print(parsed)
        with open(path, 'r', encoding='utf-8', errors='ignore') as f:
            f.seek(0, 2)
            while True:
                if line := f.readline():
                    if parsed := parse_line(line, show_all, error_only, debug, domain_filter):
                        print(parsed)
                else:
                    time.sleep(0.1)
    except KeyboardInterrupt:
        print("\nОстановка отслеживания лога." + (Style.RESET_ALL if USE_COLORS else ''))
    except Exception as e:
        print(f"{LogColors.ERROR}Ошибка при отслеживании файла: {e}{Style.RESET_ALL if USE_COLORS else ''}")

def parse_full_file(path, show_all, error_only, debug, num_lines, domain_filter):
    try:
        for line in get_last_lines(path, num_lines):
            if parsed := parse_line(line, show_all, error_only, debug, domain_filter):
                print(parsed)
    except Exception as e:
        print(f"{LogColors.ERROR}Ошибка при разборе файла: {e}{Style.RESET_ALL if USE_COLORS else ''}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Усовершенствованный анализатор логов Privoxy.',
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument("logfile", help="Путь к файлу лога Privoxy.")
    
    output_group = parser.add_mutually_exclusive_group(required=True)
    output_group.add_argument(
        "--live", type=int, nargs='?', const=1000, metavar='N', 
        help="Режим отслеживания (tail -f).\nВыводит последние N строк, затем ждет новые.\n(по умолчанию: 1000)"
    )
    output_group.add_argument(
        "--full", type=int, nargs='?', const=1000, metavar='N', 
        help="Режим разбора файла.\nВыводит последние N строк и завершает работу.\n(по умолчанию: 1000)"
    )
    
    parser.add_argument("--error", action="store_true", help="Показывать только ошибки (Crunch и HTTP >= 400).")
    parser.add_argument("--all", action="store_true", help="Показывать все записи, включая успешные HTTP 200.")
    parser.add_argument("--debug", action="store_true", help="Включить режим отладки для вывода подробной информации о парсинге.")
    parser.add_argument("--filter", help="Файл с доменами для исключения из вывода (по одному на строку).")
    parser.add_argument("--no-color", action="store_true", help="Отключить цветной вывод.")
    
    args = parser.parse_args()

    if not os.path.exists(args.logfile):
        print(f"{LogColors.ERROR}Ошибка: Файл не найден: {args.logfile}{Style.RESET_ALL}")
        exit(1)

    enable_colors(not args.no_color)
    domain_filter = DomainFilter(args.filter) if args.filter else None

    if args.live is not None:
        num_lines = args.live if args.live > 0 else 1000
        follow(args.logfile, args.all, args.error, args.debug, num_lines, domain_filter)
    elif args.full is not None:
        num_lines = args.full if args.full > 0 else 1000
        parse_full_file(args.logfile, args.all, args.error, args.debug, num_lines, domain_filter)