majuna/app/terraform/block_external.py

import logging
from collections import defaultdict

import requests

from app import app
from app.terraform.block_mirror import BlockMirrorAutomation


def _trim_prefix(s: str, prefix: str) -> str:
    if s.startswith(prefix):
        return s[len(prefix) :]
    return s


def trim_http_https(s: str) -> str:
    """
    Return the string with "http://" or "https://" removed from the start of the string if present.

    :param s: String to modify.
    :return: Modified string.
    """
    return _trim_prefix(_trim_prefix(s, "https://"), "http://")


class BlockExternalAutomation(BlockMirrorAutomation):
    """
    Automation task to import proxy reachability results from external source.
    """

    short_name = "block_external"
    description = "Import proxy reachability results from external source"

    _content: bytes

    def fetch(self) -> None:
        user_agent = {"User-agent": "BypassCensorship/1.0"}
        check_urls_config = app.config.get("EXTERNAL_CHECK_URL", [])

        if isinstance(check_urls_config, dict):
            # Config is already a dictionary, use as is.
            check_urls = check_urls_config
        elif isinstance(check_urls_config, list):
            # Convert list of strings to a dictionary with "external_N" keys.
            check_urls = {
                f"external_{i}": url for i, url in enumerate(check_urls_config)
            }
        elif isinstance(check_urls_config, str):
            # Single string, convert to a dictionary with key "external".
            check_urls = {"external": check_urls_config}
        else:
            # Fallback if the config item is neither dict, list, nor string.
            check_urls = {}
        for source, check_url in check_urls.items():
            if self._data is None:
                self._data = defaultdict(list)
            self._data[source].extend(
                requests.get(check_url, headers=user_agent, timeout=30).json()
            )

    def parse(self) -> None:
        for source, patterns in self._data.items():
            self.patterns[source].extend(
                ["https://" + trim_http_https(pattern) for pattern in patterns]
            )
        logging.debug("Found URLs: %s", self.patterns)