httphound/httphound/proxy.py

"""
Manages all proxy-related stuff
"""
import logging
import os
import subprocess
import time

from typing import Dict, Any
from dataclasses import dataclass, field
from jinja2 import Template

from .backend import BackendConfig

logger = logging.getLogger(__name__)


@dataclass
class ProxyConfig:
    """Reverse proxy configuration"""
    binary_path: str = "/usr/sbin/haproxy"
    template_path: str = "haproxy.cfg.tpl"
    working_dir: str = "/tmp/httphound"
    listen_addr: str = "*"
    listen_port: int = 4242
    template_vars: Dict[str, Any] = field(default_factory=dict)


class ProxyManager:
    """Manages reverse proxy"""

    def __init__(self, config: ProxyConfig):
        self.config = config
        self.process = None
        self.config_file = None

    def render_config(self, backend_config: BackendConfig) -> str:
        """Render proxy configuration from template"""
        template_vars = {
            'listen_addr': self.config.listen_addr,
            'listen_port': self.config.listen_port,
            'backend_host': backend_config.host,
            'backend_port': backend_config.port,
            **self.config.template_vars
        }

        try:
            with open(self.config.template_path, 'r', encoding="utf-8") as f:
                template_content = f.read()
        except FileNotFoundError:
            logger.error(f"Cannot finf template {self.config.template_path}")
            raise
        template = Template(template_content)
        rendered_template = template.render(**template_vars)
        logger.debug(f"Rendered template: \n{rendered_template}\n")
        return rendered_template

    def start(self, backend_config: BackendConfig):
        """Start the reverse proxy"""

        # Create working directory
        logger.debug(
            f"Creating reverese proxy working dir: {self.config.working_dir}")
        os.makedirs(self.config.working_dir, exist_ok=True)

        # Render and write config file
        try:
            config_content = self.render_config(backend_config)
        except Exception as e:
            logger.error(f"Cannot render template: {str(e)}")
            raise RuntimeError(f"Cannot render template: {str(e)}") from e

        self.config_file = os.path.join(self.config.working_dir, "haproxy.cfg")
        logger.debug(f"Writing configuration file {self.config_file}")
        with open(self.config_file, 'w', encoding="utf-8") as f:
            f.write(config_content)

        # Start proxy process
        # TODO: customize
        cmd = [self.config.binary_path,
               '-V',
               '-db',
               '-f', self.config_file,
               ]
        logger.debug(f"Running proxy cmd: {cmd}")

        try:
            self.process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                cwd=self.config.working_dir
            )

            # Give proxy time to start
            logger.debug("Waiting 0.1s for proxy to start")
            time.sleep(0.1)

            if self.process.poll() is not None:
                _, stderr = self.process.communicate()
                raise RuntimeError(f"Proxy failed to start: {stderr.decode()}")

            logger.debug(f"Proxy started with PID {self.process.pid}")

        except FileNotFoundError as e:
            raise RuntimeError(
                f"Proxy binary not found: {self.config.binary_path}") from e

    def stop(self):
        """Stop the reverse proxy"""
        logger.debug("Stopping reverse proxy")
        if self.process and self.process.poll() is None:
            self.process.terminate()
            try:
                self.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.process.kill()
                self.process.wait()
            logger.debug("Proxy stopped")

        # Cleanup config file
        if self.config_file and os.path.exists(self.config_file):
            logger.debug(f"Removing config file {self.config_file}")
            #os.remove(self.config_file)