token.py

"""BunnyCDN URL token authentication."""

import ipaddress
import urllib.parse
import time
import hmac
import hashlib
import base64
from typing import Dict, Optional

__version__ = "2.1.0"


def _b64url_no_pad(raw: bytes) -> str:
    """Base64url encode without padding."""
    return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")


def _build_parameters(
    parsed_query: str,
    *,
    ignore_params: bool,
    path_allowed: str,
) -> Dict[str, str]:
    if ignore_params:
        params: Dict[str, str] = {"token_ignore_params": "true"}
    else:
        raw = urllib.parse.parse_qs(parsed_query, keep_blank_values=True)
        params = {}
        for key, values in raw.items():
            if len(values) > 1:
                raise ValueError(
                    f"Multi-valued query parameter {key!r} is not supported"
                )
            params[key] = values[0]

    if path_allowed:
        params["token_path"] = path_allowed

    return dict(sorted(params.items()))


def sign_url(
    url: str,
    security_key: str,
    expiration_time: int = 86400,
    user_ip: str = "",
    is_directory: bool = True,
    path_allowed: str = "",
    countries_allowed: str = "",
    countries_blocked: str = "",
    ignore_params: bool = False,
    expires_at: Optional[int] = None,
    speed_limit: int = 0,
) -> str:
    """
    Generate a signed BunnyCDN URL.

    Args:
        url:               CDN URL without trailing '/'.
                           e.g. http://test.b-cdn.net/file.png
        security_key:      Token Authentication Key from your PullZone settings.
        expiration_time:   Token validity in seconds (default 86400 / 24 h).
                           Ignored when expires_at is set.
        user_ip:           Optional - lock the token to this IP. Must be a
                           valid IPv4 or IPv6 address when supplied.
        is_directory:      True  → token embedded in path  (/bcdn_token=...)
                           False → token in query string   (?token=...)
        path_allowed:      Optional path override for the signature scope.
        countries_allowed: Comma-separated allow-list (e.g. "CA,US,TH").
        countries_blocked: Comma-separated block-list.
        ignore_params:     If True, query params are excluded from validation.
        expires_at:        Absolute Unix timestamp for expiration. When set,
                           overrides expiration_time.

    Raises:
        ValueError: On empty/missing security_key, negative expiration,
                    multi-valued query parameters, or unparseable user_ip.
    """

    if not security_key:
        raise ValueError("security_key must not be empty")
    if expiration_time < 0:
        raise ValueError("expiration_time must be non-negative")

    parsed = urllib.parse.urlparse(url)
    query_params = urllib.parse.parse_qs(parsed.query, keep_blank_values=True)

    if countries_allowed:
        query_params["token_countries"] = [countries_allowed]
    if countries_blocked:
        query_params["token_countries_blocked"] = [countries_blocked]
    if speed_limit > 0:
        query_params["limit"] = [str(speed_limit)]

    new_query = urllib.parse.urlencode(query_params, doseq=True)
    parsed = parsed._replace(query=new_query)

    if expires_at is not None:
        expires = str(expires_at)
    else:
        expires = str(int(time.time()) + expiration_time)

    params = _build_parameters(
        parsed.query,
        ignore_params=ignore_params,
        path_allowed=path_allowed,
    )

    signature_path = path_allowed if path_allowed else parsed.path

    signing_data = "&".join(f"{k}={v}" for k, v in params.items())
    url_data = "&".join(
        f"{k}={urllib.parse.quote(v, safe='')}" for k, v in params.items()
    )

    if user_ip:
        ip_segment = ipaddress.ip_address(user_ip).packed
        if len(ip_segment) == 16:
            # Mask IPv6 to the /64 prefix: keep the first 8 bytes (network
            # portion), zero the last 8 (interface identifier). IPv4 is left unchanged.
            ip_segment = ip_segment[:8] + b"\x00" * 8
        flags_prefix = "1-"
    else:
        ip_segment = b""
        flags_prefix = ""

    message = (
        signature_path.encode("utf-8")
        + expires.encode("utf-8")
        + ip_segment
        + signing_data.encode("utf-8")
    )
    digest = hmac.new(
        security_key.encode("utf-8"),
        message,
        hashlib.sha256,
    ).digest()
    token = "HS256-" + flags_prefix + _b64url_no_pad(digest)

    base = f"{parsed.scheme}://{parsed.netloc}"
    tail = f"&{url_data}" if url_data else ""

    if is_directory:
        return f"{base}/bcdn_token={token}{tail}&expires={expires}{parsed.path}"
    else:
        return f"{base}{parsed.path}?token={token}{tail}&expires={expires}"