feat: add convenience functions, update Docker setup, and integrate Webpack for building WebRTC internals exporter

main.py

@@ -4,11 +4,13 @@ import time
 import socket
 import logging
 import os
+import argparse
 from yaspin import yaspin
 from functools import partial
 from http.server import HTTPServer
 from utils.PostHandler import Handler
 from utils.ColoredFormatter import ColoredFormatter
+from utils.Convenience import *
 from bs4 import BeautifulSoup as bs
 from selenium import webdriver
 from selenium.webdriver.chrome.options import Options
@@ -18,16 +20,29 @@ from selenium.webdriver.support.wait import WebDriverWait
 from selenium.webdriver.support import expected_conditions as ec
 
 logger = logging.getLogger(__name__)
+args = None
 
 def setupLogger():
     logging_format = "[%(asctime)s] (%(levelname)s) %(module)s - %(funcName)s: %(message)s"
-    logging.basicConfig(level=logging.INFO, format=logging_format)
+    logging.basicConfig(level=firstValid(args.log_level, os.getenv('LOG_LEVEL'), default='INFO'), format=logging_format) # type: ignore
     (logger := logging.getLogger(__name__)).setLevel(logging.INFO)
     logger.propagate = False
     (logger_handler := logging.StreamHandler()).setFormatter(
         ColoredFormatter(fmt=logging_format)
     )
     logger.addHandler(logger_handler)
+    
+
+def setupArgParser():
+    parser = argparse.ArgumentParser(description='Collector for PeerTube stats.')
+    parser.add_argument('-u', '--url', type=str, help='URL of the video to collect stats for.')
+    parser.add_argument('--socket-url', type=str, help='URL of the socket to send the stats to. Default: localhost')
+    parser.add_argument('--socket-port', type=int, help='Port of the socket to send the stats to. Default: 8094')
+    parser.add_argument('--hub-url', type=str, help='URL of the Selenium hub to connect to. If not provided, local Chrome driver will be used.')
+    parser.add_argument('--webrtc-internals-path', type=str, help='Path to the WebRTC internals extension.')
+    parser.add_argument('--log-level', type=str, help='Log level to use. Default: INFO')
+    
+    return parser
 
 def interrupt_handler(signum, driver: webdriver.Remote):
     logger.info(f'Handling signal {signum} ({signal.Signals(signum).name}).')
@@ -36,37 +51,40 @@ def interrupt_handler(signum, driver: webdriver.Remote):
     raise SystemExit
 
 @yaspin()
-def setupChromeDriver():
+def setupChromeDriver(command_executor: str | None, webrtc_internals_path: str) -> webdriver.Remote | webdriver.Chrome:
     logger.log(logging.INFO, 'Setting up Chrome driver.')
     chrome_options = Options()
     #chrome_options.add_argument("--headless")
     chrome_options.add_argument("--no-sandbox")
     chrome_options.add_argument("--mute-audio")
     chrome_options.add_argument("--window-size=1280,720")
-    #chrome_options.add_argument("--disable-dev-shm-usage")
     chrome_options.add_argument("--no-default-browser-check")
     chrome_options.add_argument("--disable-features=WebRtcHideLocalIpsWithMdns")
-    #chrome_options.add_argument(f"--load-extension={os.path.abspath(os.path.join(os.path.dirname(__file__), 'webrtc-internals-exporter'))}")
-    chrome_options.add_argument("--load-extension=/tmp/webrtc-internals-exporter")
+    chrome_options.add_argument(f"--load-extension={webrtc_internals_path}")
     chrome_options.add_experimental_option('prefs', {'intl.accept_languages': 'en,en_US'})
-        
-    #driver = webdriver.Chrome(options=chrome_options)
-    driver = webdriver.Remote(command_executor='http://host.docker.internal:4444', options=chrome_options)
+
+    if command_executor is not None:
+        driver = webdriver.Remote(command_executor=command_executor, options=chrome_options)
+        logger.warning(f'Using Selenium hub at {command_executor}.')
+    else:
+        driver = webdriver.Chrome(options=chrome_options)
+        logger.warning('No Selenium hub URL provided, using local Chrome driver.')
+    
     logger.log(logging.INFO, 'Chrome driver setup complete.')
     
     return driver
 
-def saveStats(stats: list):
+def saveStats(stats: list, socket_url: str, socket_port: int):
     try:
         sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
         logger.log(logging.DEBUG, f'Saving stats: {json.dumps(stats, indent=4)}')
-        sock.sendto(json.dumps(stats).encode(), ('telegraf', 8094))
+        sock.sendto(json.dumps(stats).encode(), (socket_url, socket_port))
         sock.close()
         logger.log(logging.DEBUG, 'Sent stats to socket.')
     except socket.error as e:
         logger.error(f'Got socket error: {e}')
 
-def downloadStats(driver: webdriver.Chrome, peersDict: dict):
+def downloadStats(driver: webdriver.Remote | webdriver.Chrome, peersDict: dict, socket_url: str, socket_port: int):
     html = driver.find_element(By.CLASS_NAME ,'vjs-stats-list').get_attribute('innerHTML')
     if html is not None:
         htmlBS = bs(html, 'html.parser')
@@ -157,7 +175,7 @@ def downloadStats(driver: webdriver.Chrome, peersDict: dict):
         'session': driver.session_id
     }
 
-    saveStats([stats])
+    saveStats([stats], socket_url, socket_port)
 
 def convert_to_bytes(down, downUnit):
     return float(down) * (1024 ** {'B': 0, 'KB': 1, 'MB': 2, 'GB': 3}[downUnit])
@@ -188,20 +206,31 @@ def setupStats(driver: webdriver.Remote, url: str):
     
     return driver
 
-if __name__ == '__main__':    
+if __name__ == '__main__':
+    args = setupArgParser().parse_args()
+    
     setupLogger()
     
-    driver = setupChromeDriver()
+    command_executor = firstValid(args.hub_url, os.getenv('HUB_URL'), default=None)
+    webrtc_internals_path = firstValid(
+        args.webrtc_internals_path,
+        os.getenv('WEBRTC_INTERNALS_PATH'), 
+        default=os.path.abspath(os.path.join(os.path.dirname(__file__), 'webrtc-internals-exporter'))
+        )
+    
+    driver = setupChromeDriver(command_executor, webrtc_internals_path)
 
     signal.signal(signal.SIGINT, lambda signum, frame: interrupt_handler(signum, driver))
     
-    url = os.getenv('VIDEO_URL')
+    url = firstValid(args.url, os.getenv('VIDEO_URL'), default=None)
     if url is None:
-        logger.error('VIDEO_URL environment variable is not set.')
+        logger.error('VIDEO_URL environment variable or --url argument is required.')
         raise SystemExit(1)
     
     setupStats(driver, url)
     
+    socket_url = firstValid(args.socket_url, os.getenv('SOCKET_URL'), default='localhost')
+    socket_port = firstValid(args.socket_port, os.getenv('SOCKET_PORT'), default=8094)
     logger.info('Starting server collector.')
-    httpd = HTTPServer(('', 9092), partial(Handler, downloadStats, driver, logger))
+    httpd = HTTPServer(('', 9092), partial(Handler, downloadStats, driver, logger, socket_url, socket_port))
     httpd.serve_forever()