mirror of
https://github.com/Dispatcharr/Dispatcharr.git
synced 2026-07-22 09:37:57 +00:00
225 lines
8.2 KiB
Python
225 lines
8.2 KiB
Python
"""
|
|
Transport Stream (TS) Proxy Server
|
|
Handles live TS stream proxying with support for:
|
|
- Stream switching
|
|
- Buffer management
|
|
- Multiple client connections
|
|
- Connection state tracking
|
|
"""
|
|
|
|
import requests
|
|
import threading
|
|
import logging
|
|
from collections import deque
|
|
import time
|
|
from typing import Optional, Set, Deque, Dict
|
|
from apps.proxy.config import TSConfig as Config
|
|
|
|
class StreamManager:
|
|
"""Manages TS stream state and connection handling"""
|
|
|
|
def __init__(self, initial_url: str, channel_id: str, user_agent: Optional[str] = None):
|
|
self.current_url: str = initial_url
|
|
self.channel_id: str = channel_id
|
|
self.user_agent: str = user_agent or Config.DEFAULT_USER_AGENT
|
|
self.url_changed: threading.Event = threading.Event()
|
|
self.running: bool = True
|
|
self.session: requests.Session = self._create_session()
|
|
self.connected: bool = False
|
|
self.retry_count: int = 0
|
|
logging.info(f"Initialized stream manager for channel {channel_id}")
|
|
|
|
def _create_session(self) -> requests.Session:
|
|
"""Create and configure requests session"""
|
|
session = requests.Session()
|
|
session.headers.update({
|
|
'User-Agent': self.user_agent,
|
|
'Connection': 'keep-alive'
|
|
})
|
|
return session
|
|
|
|
def update_url(self, new_url: str) -> bool:
|
|
"""Update stream URL and signal connection change"""
|
|
if new_url != self.current_url:
|
|
logging.info(f"Stream switch initiated: {self.current_url} -> {new_url}")
|
|
self.current_url = new_url
|
|
self.connected = False
|
|
self.url_changed.set()
|
|
return True
|
|
return False
|
|
|
|
def should_retry(self) -> bool:
|
|
"""Check if connection retry is allowed"""
|
|
return self.retry_count < Config.MAX_RETRIES
|
|
|
|
def stop(self) -> None:
|
|
"""Clean shutdown of stream manager"""
|
|
self.running = False
|
|
if self.session:
|
|
self.session.close()
|
|
|
|
class StreamBuffer:
|
|
"""Manages stream data buffering"""
|
|
|
|
def __init__(self):
|
|
self.buffer: Deque[bytes] = deque(maxlen=Config.BUFFER_SIZE)
|
|
self.lock: threading.Lock = threading.Lock()
|
|
self.index: int = 0
|
|
|
|
class ClientManager:
|
|
"""Manages active client connections"""
|
|
|
|
def __init__(self):
|
|
self.active_clients: Set[int] = set()
|
|
self.lock: threading.Lock = threading.Lock()
|
|
|
|
def add_client(self, client_id: int) -> None:
|
|
"""Add new client connection"""
|
|
with self.lock:
|
|
self.active_clients.add(client_id)
|
|
logging.info(f"New client connected: {client_id} (total: {len(self.active_clients)})")
|
|
|
|
def remove_client(self, client_id: int) -> int:
|
|
"""Remove client and return remaining count"""
|
|
with self.lock:
|
|
self.active_clients.remove(client_id)
|
|
remaining = len(self.active_clients)
|
|
logging.info(f"Client disconnected: {client_id} (remaining: {remaining})")
|
|
return remaining
|
|
|
|
class StreamFetcher:
|
|
"""Handles stream data fetching"""
|
|
|
|
def __init__(self, manager: StreamManager, buffer: StreamBuffer):
|
|
self.manager = manager
|
|
self.buffer = buffer
|
|
|
|
def fetch_loop(self) -> None:
|
|
"""Main fetch loop for stream data"""
|
|
while self.manager.running:
|
|
try:
|
|
if not self._handle_connection():
|
|
continue
|
|
|
|
with self.manager.session.get(self.manager.current_url, stream=True) as response:
|
|
if response.status_code == 200:
|
|
self._handle_successful_connection()
|
|
self._process_stream(response)
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
self._handle_connection_error(e)
|
|
|
|
def _handle_connection(self) -> bool:
|
|
"""Handle connection state and retries"""
|
|
if not self.manager.connected:
|
|
if not self.manager.should_retry():
|
|
logging.error(f"Failed to connect after {Config.MAX_RETRIES} attempts")
|
|
return False
|
|
|
|
if not self.manager.running:
|
|
return False
|
|
|
|
self.manager.retry_count += 1
|
|
logging.info(f"Connecting to stream: {self.manager.current_url} "
|
|
f"(attempt {self.manager.retry_count}/{Config.MAX_RETRIES})")
|
|
return True
|
|
|
|
def _handle_successful_connection(self) -> None:
|
|
"""Handle successful stream connection"""
|
|
if not self.manager.connected:
|
|
logging.info("Stream connected successfully")
|
|
self.manager.connected = True
|
|
self.manager.retry_count = 0
|
|
|
|
def _process_stream(self, response: requests.Response) -> None:
|
|
"""Process incoming stream data"""
|
|
for chunk in response.iter_content(chunk_size=Config.CHUNK_SIZE):
|
|
if not self.manager.running:
|
|
logging.info("Stream fetch stopped - shutting down")
|
|
return
|
|
|
|
if chunk:
|
|
if self.manager.url_changed.is_set():
|
|
logging.info("Stream switch in progress, closing connection")
|
|
self.manager.url_changed.clear()
|
|
break
|
|
|
|
with self.buffer.lock:
|
|
self.buffer.buffer.append(chunk)
|
|
self.buffer.index += 1
|
|
|
|
def _handle_connection_error(self, error: Exception) -> None:
|
|
"""Handle stream connection errors"""
|
|
logging.error(f"Stream connection error: {error}")
|
|
self.manager.connected = False
|
|
|
|
if not self.manager.running:
|
|
return
|
|
|
|
logging.info(f"Attempting to reconnect in {Config.RECONNECT_DELAY} seconds...")
|
|
if not wait_for_running(self.manager, Config.RECONNECT_DELAY):
|
|
return
|
|
|
|
def wait_for_running(manager: StreamManager, delay: float) -> bool:
|
|
"""Wait while checking manager running state"""
|
|
start = time.time()
|
|
while time.time() - start < delay:
|
|
if not manager.running:
|
|
return False
|
|
threading.Event().wait(0.1)
|
|
return True
|
|
|
|
class ProxyServer:
|
|
"""Manages TS proxy server instance"""
|
|
|
|
def __init__(self, user_agent: Optional[str] = None):
|
|
self.stream_managers: Dict[str, StreamManager] = {}
|
|
self.stream_buffers: Dict[str, StreamBuffer] = {}
|
|
self.client_managers: Dict[str, ClientManager] = {}
|
|
self.fetch_threads: Dict[str, threading.Thread] = {}
|
|
self.user_agent: str = user_agent or Config.DEFAULT_USER_AGENT
|
|
|
|
def initialize_channel(self, url: str, channel_id: str) -> None:
|
|
"""Initialize a new channel stream"""
|
|
if channel_id in self.stream_managers:
|
|
self.stop_channel(channel_id)
|
|
|
|
self.stream_managers[channel_id] = StreamManager(
|
|
url,
|
|
channel_id,
|
|
user_agent=self.user_agent
|
|
)
|
|
self.stream_buffers[channel_id] = StreamBuffer()
|
|
self.client_managers[channel_id] = ClientManager()
|
|
|
|
fetcher = StreamFetcher(
|
|
self.stream_managers[channel_id],
|
|
self.stream_buffers[channel_id]
|
|
)
|
|
|
|
self.fetch_threads[channel_id] = threading.Thread(
|
|
target=fetcher.fetch_loop,
|
|
name=f"StreamFetcher-{channel_id}",
|
|
daemon=True
|
|
)
|
|
self.fetch_threads[channel_id].start()
|
|
logging.info(f"Initialized channel {channel_id} with URL {url}")
|
|
|
|
def stop_channel(self, channel_id: str) -> None:
|
|
"""Stop and cleanup a channel"""
|
|
if channel_id in self.stream_managers:
|
|
self.stream_managers[channel_id].stop()
|
|
if channel_id in self.fetch_threads:
|
|
self.fetch_threads[channel_id].join(timeout=5)
|
|
self._cleanup_channel(channel_id)
|
|
|
|
def _cleanup_channel(self, channel_id: str) -> None:
|
|
"""Remove channel resources"""
|
|
for collection in [self.stream_managers, self.stream_buffers,
|
|
self.client_managers, self.fetch_threads]:
|
|
collection.pop(channel_id, None)
|
|
|
|
def shutdown(self) -> None:
|
|
"""Stop all channels and cleanup"""
|
|
for channel_id in list(self.stream_managers.keys()):
|
|
self.stop_channel(channel_id)
|