ResolveShortenedURL

Source

from abc import ABCMeta
from typing import NamedTuple

import demistomock as demisto  # noqa: F401
import urllib3
from CommonServerPython import *  # noqa: F401
from requests import Response

urllib3.disable_warnings()  # Disable insecure warnings

DEFAULT_SERVICE = "Built-In"
DEFAULT_REDIRECT_LIMIT = "0"


class URLUnshorteningData(NamedTuple):
    """
    A tuple containing data for unshortend URLs.

    Attributes:
        original_url (str): The original URL.
        resolved_url (str): The resolved URL.
        service_name (str): The name of the service used to resolve the URL.
        redirect_history (list): A list of URLs that were redirected to get to the resolved URL.
        raw_data (dict | list[dict] | None, optional): The raw data returned by the service. None if not available.
        encountered_error (bool, optional): Whether an error was encountered while resolving the URL. Defaults to False.
        api_usage (int | None, optional): The API usage count for the current IP. None if not relevant to the service.
        api_rate_limit (int | None, optional): The maximum number of API calls allowed by the service
          for a defined period of time. None if not relevant to the service.
    """

    original_url: str
    resolved_url: str
    service_name: str
    redirect_history: list[str]
    raw_data: dict | list[dict] | None = None
    encountered_error: bool = False
    api_usage: int | None = None
    api_rate_limit: int | None = None

    def to_context_dict(self) -> dict:
        """
        Converts the data to a dictionary that will be used as the context data.
        Adds recursion data only if relevant.

        Note:
            We subtract 1 from RedirectCount because the original URL is included in the recursion history.

        Returns:
            dict: A dictionary containing the data in context format.
        """
        data = {
            "OriginalURL": self.original_url,
            "ResolvedURL": self.resolved_url,
            "ServiceName": self.service_name,
            "RedirectCount": len(self.redirect_history) - 1,
            "RedirectHistory": self.redirect_history,
            "EncounteredError": self.encountered_error,
        }

        if self.api_usage is not None:
            data["APIUsageCount"] = self.api_usage

        if self.api_rate_limit is not None:
            data["APIRateLimit"] = self.api_rate_limit

        return data

    def to_hr_dict(self) -> dict:
        """
        Converts the data to a dictionary that will be used as the human-readable data.

        Returns:
            dict: A dictionary containing the data in human-readable format.
        """
        api_usage_hr: str | None = None

        if self.api_usage is not None:
            api_usage_hr = f"{self.api_usage}"

            if self.api_rate_limit is not None:
                api_usage_hr += f"/{self.api_rate_limit}"

        return {
            "Original URL": self.original_url,
            "Resolved URL": self.resolved_url,
            "Service Used": self.service_name,
            "Redirect History": self.redirect_history if len(self.redirect_history) > 1 else None,
            "API Usage Count": api_usage_hr if api_usage_hr is not None else None,
        }


class URLUnshortingService(BaseClient, metaclass=ABCMeta):
    """
    An abstract base class for URL unshorteners.

    Note:
        To add a new service, create a new class that inherits from this class, and implements the `resolve_url` method.
        The class attribute `service_name` must match the name used in the under `service` on the YAML file.
        Once created new service should be automatically detected and used.

    Attributes:
        base_url (str): The base URL of the service that will be used for sending requests.
        service_name (str): The name of the service.
        redirect_limit (int | None): The maximum number of redirects to follow. None if no limit.
    """

    base_url: str
    service_name: str
    service_rate_limit: int | None = None

    def __init__(self, redirect_limit: int | None = None, **kwargs):
        super().__init__(base_url=self.base_url, **kwargs)
        self.redirect_limit = redirect_limit if redirect_limit is not None else 0

    def hit_redirect_limit(self, redirect_history: list) -> bool:
        """
        Checks whether the redirect limit has been reached.

        Args:
            redirect_history (list): The redirect history of the URL.

        Returns:
            bool: True if the redirect limit has been reached, False otherwise.
        """
        if self.redirect_limit == 0:
            return False

        return len(redirect_history) >= self.redirect_limit

    @staticmethod
    def find_matching_service(service_name: str) -> type["URLUnshortingService"]:
        """
        Finds a matching service class by name.

        Args:
            service_name (str): The service name to find (has to match the service_name class attribute).

        Returns:
            Type[URLUnshortingService]: A subclass of URLUnshortingService that matches the service name.
        """
        for service_class in URLUnshortingService.__subclasses__():
            if service_class.service_name.casefold() == service_name.casefold():
                return service_class

        raise ValueError(f'No matching service was found for: "{service_name}".')

    @abstractmethod
    def resolve_url(self, url: str) -> URLUnshorteningData:  # pragma: no cover
        """
        Resolve a shortened URL.

        Args:
            url (str): The URL to resolve.

        Returns:
            URLUnshorteningData: A NamedTuple containing the data for the resolved URL.
        """


class BuiltInShortener(URLUnshortingService):
    """A class for unshortening URLs using Python requests."""

    base_url = ""
    service_name = "Built-In"

    def resolve_url(self, url: str) -> URLUnshorteningData:
        encountered_error = False
        original_url = url
        url = url if url.startswith("http") else f"https://{url}"

        try:
            response: Response = self._http_request(
                method="GET",
                full_url=url,
                resp_type="response",
                allow_redirects=False,
            )

        except Exception:
            encountered_error = True

        redirect_history: list[str] = []

        while not encountered_error and (response.is_redirect and (not self.hit_redirect_limit(redirect_history))):
            url = response.headers["location"]
            redirect_history.append(url)

            try:
                response = self._http_request(
                    method="GET",
                    full_url=url,
                    resp_type="response",
                    allow_redirects=False,
                )

            except Exception:
                if not self.hit_redirect_limit(redirect_history):
                    encountered_error = True
                break

        return URLUnshorteningData(
            original_url=original_url,
            resolved_url=url,
            service_name=self.service_name,
            redirect_history=[original_url] + redirect_history,
            encountered_error=encountered_error,
        )


def unshorten_url(
    service_name: str, url: str, redirect_limit: int, use_system_proxy: bool = False, session_verify: bool = True
) -> CommandResults:
    """
    Unshorten a shortened URL.

    Args:
        service_name (str): The service to use for unshortening.
        url (str): The URL to un-shorten.
        use_system_proxy (bool): Whether to use the system proxy.
        session_verify (bool): Whether to verify the SSL certificate of the request.
        redirect_limit (int): A maximum number of recursions to run. Use 0 for unlimited.
    """
    error_message = (
        "There was an error while attempting to unshorten the final URL in the redirect chain.\n"
        "It is possible that the unshortening process was not fully completed.\n\n"
    )

    service_class = URLUnshortingService.find_matching_service(service_name=service_name)
    service_instance = service_class(  # pylint: disable=E0110
        redirect_limit=redirect_limit,  # pylint: disable=E0110
        proxy=use_system_proxy,
        verify=session_verify,
    )
    returned_data = service_instance.resolve_url(url=url)

    readable_output = ""

    if returned_data.encountered_error:
        readable_output += error_message

    readable_output += tableToMarkdown(
        name="URL Unshortening Results",
        t=returned_data.to_hr_dict(),
        headers=list(returned_data.to_hr_dict().keys()),
        removeNull=True,
    )

    return CommandResults(
        outputs_prefix="ResolveShortenedURL",
        outputs_key_field="OriginalURL",
        outputs=returned_data.to_context_dict(),
        readable_output=readable_output,
        indicator=Common.URL(
            url=returned_data.resolved_url,
            dbot_score=Common.DBotScore(
                indicator=returned_data.resolved_url,
                indicator_type=DBotScoreType.URL,
                integration_name="ResolveShortenedURL",
                score=Common.DBotScore.NONE,
            ),
        ),
        raw_response=returned_data.raw_data,
    )


def main():  # pragma: no cover
    args = demisto.args()

    try:
        url: str = args["url"]
        use_system_proxy = argToBoolean(args.get("use_system_proxy", "False"))
        redirect_limit = arg_to_number(args.get("redirect_limit", DEFAULT_REDIRECT_LIMIT))

        # `arg_to_number` returns `None` if int conversion was unsuccessful.
        if redirect_limit is None:
            raise ValueError("'redirect_limit' must be a natural number.")

        session_verify = not argToBoolean(args.get("insecure", "False"))

        result = unshorten_url(
            service_name=DEFAULT_SERVICE,
            url=url,
            use_system_proxy=use_system_proxy,
            redirect_limit=redirect_limit,
            session_verify=session_verify,
        )
        return_results(result)

    except Exception as e:
        return_error(f"Error: {e!s}")


if __name__ in ("__main__", "__builtin__", "builtins"):  # pragma: no cover
    main()

README