227 lines
6.9 KiB
Python
227 lines
6.9 KiB
Python
"""Proxy pool management for web crawling.
|
|
|
|
Provides multi-proxy configuration, round-robin rotation, and
|
|
backward compatibility with single proxy environment variables.
|
|
|
|
Configuration file (~/.kb_proxies.yaml):
|
|
proxies:
|
|
- "http://127.0.0.1:7890"
|
|
- "http://user:pass@proxy.example.com:8080"
|
|
- "socks5://127.0.0.1:1080"
|
|
cooldown_seconds: 300
|
|
max_fail_count: 3
|
|
"""
|
|
|
|
import os
|
|
import logging
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import yaml
|
|
from pydantic import BaseModel, Field
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ProxyEntry(BaseModel):
|
|
"""Single proxy configuration and state.
|
|
|
|
Tracks proxy URL, failure count, and availability status.
|
|
|
|
Attributes:
|
|
url: Complete proxy URL including authentication
|
|
fail_count: Consecutive failure count
|
|
last_fail_time: Timestamp of last failure
|
|
is_available: Whether proxy is currently available
|
|
"""
|
|
|
|
url: str = Field(..., description="Complete proxy URL including auth")
|
|
fail_count: int = Field(default=0, description="Consecutive failure count")
|
|
last_fail_time: Optional[datetime] = Field(
|
|
default=None, description="Timestamp of last failure"
|
|
)
|
|
is_available: bool = Field(default=True, description="Whether proxy is available")
|
|
|
|
def record_success(self) -> None:
|
|
"""Record successful request, reset failure count."""
|
|
self.fail_count = 0
|
|
self.is_available = True
|
|
|
|
def record_failure(self, max_fail_count: int = 3) -> None:
|
|
"""Record failed request.
|
|
|
|
Args:
|
|
max_fail_count: Threshold for marking proxy unavailable
|
|
"""
|
|
self.fail_count += 1
|
|
self.last_fail_time = datetime.now()
|
|
if self.fail_count >= max_fail_count:
|
|
self.is_available = False
|
|
logger.warning(f"Proxy marked as unavailable: {self.url[:30]}...")
|
|
|
|
def check_recovery(self, cooldown_seconds: int = 300) -> bool:
|
|
"""Check if proxy can recover from cooldown.
|
|
|
|
Lazy check - only evaluates when called.
|
|
|
|
Args:
|
|
cooldown_seconds: Minimum time before recovery
|
|
|
|
Returns:
|
|
True if proxy is available or has recovered
|
|
"""
|
|
if self.is_available:
|
|
return True
|
|
if self.last_fail_time is None:
|
|
return True
|
|
|
|
elapsed = (datetime.now() - self.last_fail_time).total_seconds()
|
|
if elapsed >= cooldown_seconds:
|
|
self.is_available = True
|
|
self.fail_count = 0
|
|
logger.info(f"Proxy recovered after {elapsed:.0f}s cooldown")
|
|
return True
|
|
return False
|
|
|
|
|
|
class ProxyPoolConfig(BaseModel):
|
|
"""Proxy pool configuration.
|
|
|
|
Attributes:
|
|
proxies: List of proxy URLs
|
|
cooldown_seconds: Time before failed proxy can recover
|
|
max_fail_count: Consecutive failures before marking unavailable
|
|
"""
|
|
|
|
proxies: list[str] = Field(default_factory=list, description="Proxy URL list")
|
|
cooldown_seconds: int = Field(default=300, description="Cooldown time in seconds")
|
|
max_fail_count: int = Field(default=3, description="Max consecutive failures")
|
|
|
|
|
|
class ProxyPool:
|
|
"""Proxy pool manager.
|
|
|
|
Manages multiple proxies with:
|
|
- YAML file or environment variable loading
|
|
- Round-robin rotation
|
|
- Passive health detection
|
|
- Cooldown recovery
|
|
|
|
Usage:
|
|
pool = ProxyPool.from_yaml(Path.home() / ".kb_proxies.yaml")
|
|
if pool.has_proxies():
|
|
proxy = pool.get_next_proxy()
|
|
"""
|
|
|
|
def __init__(self, config: ProxyPoolConfig):
|
|
"""Initialize proxy pool.
|
|
|
|
Args:
|
|
config: Proxy pool configuration
|
|
"""
|
|
self.config = config
|
|
self._entries: list[ProxyEntry] = [
|
|
ProxyEntry(url=url) for url in config.proxies
|
|
]
|
|
self._current_index: int = 0
|
|
|
|
@classmethod
|
|
def from_yaml(cls, path: Path) -> "ProxyPool":
|
|
"""Load proxy pool from YAML file.
|
|
|
|
Args:
|
|
path: Path to YAML configuration file
|
|
|
|
Returns:
|
|
ProxyPool instance (empty if file doesn't exist)
|
|
"""
|
|
if not path.exists():
|
|
logger.debug(f"Proxy config file not found: {path}")
|
|
return cls(ProxyPoolConfig())
|
|
|
|
with open(path) as f:
|
|
data = yaml.safe_load(f) or {}
|
|
|
|
proxies = data.get("proxies", [])
|
|
cooldown_seconds = data.get("cooldown_seconds", 300)
|
|
max_fail_count = data.get("max_fail_count", 3)
|
|
|
|
return cls(
|
|
ProxyPoolConfig(
|
|
proxies=proxies,
|
|
cooldown_seconds=cooldown_seconds,
|
|
max_fail_count=max_fail_count,
|
|
)
|
|
)
|
|
|
|
@classmethod
|
|
def from_env_fallback(cls) -> "ProxyPool":
|
|
"""Load single proxy from environment variable.
|
|
|
|
Checks HTTPS_PROXY first, then HTTP_PROXY.
|
|
Provides backward compatibility with single proxy mode.
|
|
|
|
Returns:
|
|
ProxyPool instance with single proxy or empty
|
|
"""
|
|
proxy_url = os.environ.get("HTTPS_PROXY") or os.environ.get("HTTP_PROXY")
|
|
if proxy_url:
|
|
logger.debug(f"Using proxy from environment: {proxy_url[:30]}...")
|
|
return cls(ProxyPoolConfig(proxies=[proxy_url]))
|
|
return cls(ProxyPoolConfig())
|
|
|
|
def get_next_proxy(self) -> Optional[ProxyEntry]:
|
|
"""Get next available proxy using round-robin.
|
|
|
|
Implements:
|
|
1. Lazy recovery check for current proxy
|
|
2. Skip unavailable proxies
|
|
3. Fallback to first proxy if all unavailable
|
|
|
|
Returns:
|
|
Next ProxyEntry or None if pool is empty
|
|
"""
|
|
if not self._entries:
|
|
return None
|
|
|
|
# Single proxy mode: always return the same proxy
|
|
if len(self._entries) == 1:
|
|
return self._entries[0]
|
|
|
|
# Try to find available proxy
|
|
checked: set[int] = set()
|
|
while len(checked) < len(self._entries):
|
|
entry = self._entries[self._current_index]
|
|
# Move index before any potential early returns
|
|
self._current_index = (self._current_index + 1) % len(self._entries)
|
|
checked.add(self._current_index)
|
|
|
|
# Lazy recovery check
|
|
if entry.check_recovery(self.config.cooldown_seconds):
|
|
return entry
|
|
|
|
# All proxies unavailable, fallback to first proxy
|
|
logger.warning("All proxies unavailable, falling back to first proxy")
|
|
return self._entries[0]
|
|
|
|
def has_proxies(self) -> bool:
|
|
"""Check if pool has configured proxies.
|
|
|
|
Returns:
|
|
True if pool has at least one proxy
|
|
"""
|
|
return len(self._entries) > 0
|
|
|
|
def get_available_count(self) -> int:
|
|
"""Get count of currently available proxies.
|
|
|
|
Returns:
|
|
Number of available proxies
|
|
"""
|
|
return sum(
|
|
1
|
|
for e in self._entries
|
|
if e.check_recovery(self.config.cooldown_seconds)
|
|
)
|