cmsdials 0.0.0__py3-none-any.whl

cmsdials/auth/_base.py

@@ -0,0 +1,136 @@
+import threading
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+from datetime import datetime, timedelta
+from enum import Enum
+
+from .exceptions import ImpossibleToRefreshTokenError
+
+
+class TokenState(Enum):
+    """
+    Tracks the state of a token.
+
+    FRESH: The token is valid and can be used.
+    STALE: The token is close to expired, it can be used but you also should trigger a non blocking refresh.
+    EXPIRED: The token is expired but the refresh token is valid. You must trigger a block refresh.
+    INVALID: Both the token and refresh token are expired.
+    """
+
+    FRESH = 1
+    STALE = 2
+    EXPIRED = 3
+    INVALID = 4
+
+
+class BaseCredentials(ABC):
+    """
+    Base class for all credentials.
+    """
+
+    def __init__(self):
+        self.token = None
+        self.expires_at = None
+        self.expires_in = None
+        self.refresh_expires_at = None
+        self.refresh_expires_in = None
+        self.refresh_token = None
+        self.token_type = None
+        self._stale_lock = threading.Lock()
+        self._expired_lock = threading.Lock()
+
+    @property
+    def stale(self) -> bool:
+        """
+        States if token is staled.
+        """
+        if not self.expires_at:
+            return False
+        return self.expires_at > datetime.now() > (self.expires_at - timedelta(seconds=0.1 * self.expires_in))
+
+    @property
+    def expired(self) -> bool:
+        """
+        States if token is expired.
+        """
+        if not self.expires_at:
+            return False
+        return datetime.now() >= self.expires_at
+
+    @property
+    def refresh_expired(self) -> bool:
+        """
+        States if refresh token is expired.
+        """
+        if not self.refresh_expires_at:
+            return False
+        return datetime.now() >= self.refresh_expires_at
+
+    @property
+    def token_state(self):
+        """
+        See `:obj:`TokenState`
+        """
+        # Credentials that can't expire are always treated as fresh.
+        if self.expires_at is None:
+            return TokenState.FRESH
+
+        if self.stale:
+            return TokenState.STALE
+
+        if self.expired and not self.refresh_expired:
+            return TokenState.EXPIRED
+
+        if self.expired and self.refresh_expired:
+            return TokenState.INVALID
+
+        return TokenState.FRESH
+
+    @abstractmethod
+    def apply(self, headers: Mapping[str, str]) -> None:
+        """
+        Apply the token to the authentication header.
+        """
+        raise NotImplementedError("apply must be implemented")
+
+    @abstractmethod
+    def refresh(self) -> None:
+        """
+        Performs credential-specific refresh logic.
+        """
+        raise NotImplementedError("refresh must be implemented")
+
+    def _blocking_refresh(self) -> None:
+        if self.token_state == TokenState.EXPIRED:
+            self.refresh()
+
+    def _nonblocking_refresh(self) -> None:
+        self._stale_lock.acquire()
+        self.refresh()
+        self._stale_lock.release()
+
+    def before_request(self, headers: Mapping[str, str]):
+        """
+        Performs credential-specific before request logic.
+        """
+        # If token is invalid (both token and refresh token expired): Throw error
+        # User need to interactively authenticate again
+        if self.token_state == TokenState.INVALID:
+            raise ImpossibleToRefreshTokenError
+
+        # If token is stale (token is close to expire)
+        # Re-use the token but launch a new thread to refresh the token
+        if self.token_state == TokenState.STALE and self._stale_lock.locked() is False:
+            worker = threading.Thread(target=self._nonblocking_refresh)
+            worker.start()
+
+        # If token is expired (token is expired but refresh token is not expired)
+        # Do not re-use the token, block the main thread and refresh before continuing
+        # After refresh the token is fresh
+        if self.token_state == TokenState.EXPIRED:
+            self._expired_lock.acquire()
+            self._blocking_refresh()
+            self._expired_lock.release()
+
+        # Put token in the header if token is fresh
+        self.apply(headers)

cmsdials/auth/bearer.py

@@ -0,0 +1,140 @@
+import json
+import os
+import os.path
+from collections.abc import Mapping
+from datetime import datetime
+from typing import Optional
+
+from ..utils._json import TokenDecoder, TokenEncoder
+from ..utils.logger import logger
+from ._base import BaseCredentials, TokenState
+from .client import AuthClient
+from .models import Token
+
+
+DEFAULT_CACHE_DIR = ".cache"
+DEFAULT_CACHE_FILENAME = "creds"
+
+
+class Credentials(BaseCredentials):
+    """
+    Token credentials.
+    These credential use a token to provide authorization to applications.
+    """
+
+    def __init__(
+        self,
+        token: str,
+        expires_at: datetime,
+        expires_in: int,
+        refresh_expires_at: datetime,
+        refresh_expires_in: int,
+        refresh_token: str,
+        token_type: str,
+        cache_dir: str = DEFAULT_CACHE_DIR,
+        client: Optional[AuthClient] = None,
+    ):
+        # Load parent class and overwrite token attributes
+        super().__init__()
+        self.token = token
+        self.expires_at = expires_at
+        self.expires_in = expires_in
+        self.refresh_expires_at = refresh_expires_at
+        self.refresh_expires_in = refresh_expires_in
+        self.refresh_token = refresh_token
+        self.token_type = token_type
+
+        # Setup dir for caching credentials and authentication client
+        self.cache_dir = cache_dir
+        self.client = client or AuthClient()
+
+        # If loaded token is expired, immediately refresh
+        if self.token_state == TokenState.EXPIRED:
+            self.refresh()
+
+        # If loaded token is invalid, trigger authentication
+        if self.token_state == TokenState.INVALID:
+            logger.info("Access token and refresh token are expired, triggering device authentication flow...")
+            token: Token = self.client.device_auth_flow()
+            self._set_token(token)
+
+        # Only cache if token is fresh
+        if self.token_state == TokenState.FRESH:
+            self.cache_credentials()
+
+    @property
+    def _token(self) -> Token:
+        return Token(
+            access_token=self.token,
+            expires_in=self.expires_in,
+            refresh_expires_in=self.refresh_expires_in,
+            refresh_token=self.refresh_token,
+            token_type=self.token_type,
+            expires_at=self.expires_at,
+            refresh_expires_at=self.refresh_expires_at,
+        )
+
+    @staticmethod
+    def _handle_creds_filepath(cache_dir: str = DEFAULT_CACHE_DIR):
+        cache_dir = cache_dir or DEFAULT_CACHE_DIR
+        if os.path.isabs(cache_dir) is False:
+            cache_dir = os.path.join(os.getcwd(), cache_dir)
+
+        os.makedirs(cache_dir, exist_ok=True)
+        return os.path.join(cache_dir, DEFAULT_CACHE_FILENAME)
+
+    @staticmethod
+    def from_authclient_token(
+        token: Token, cache_dir: str = DEFAULT_CACHE_DIR, client: Optional[AuthClient] = None
+    ) -> "Credentials":
+        client = client or AuthClient()
+        return Credentials(
+            token=token.access_token,
+            expires_in=token.expires_in,
+            refresh_expires_in=token.refresh_expires_in,
+            refresh_token=token.refresh_token,
+            token_type=token.token_type,
+            expires_at=token.expires_at,
+            refresh_expires_at=token.refresh_expires_at,
+            cache_dir=cache_dir,
+            client=client,
+        )
+
+    @staticmethod
+    def from_creds_file(cache_dir: str = DEFAULT_CACHE_DIR, client: Optional[AuthClient] = None) -> "Credentials":
+        client = client or AuthClient()
+        fpath = Credentials._handle_creds_filepath(cache_dir)
+        try:
+            with open(fpath) as filebuffer:
+                fcontents = json.load(filebuffer, cls=TokenDecoder)
+            token = Token(**fcontents)
+        except FileNotFoundError:
+            logger.info("Credentials file not found, triggering device authentication flow...")
+            token: Token = client.device_auth_flow()
+
+        return Credentials.from_authclient_token(token, cache_dir, client)
+
+    def _set_token(self, token: Token):
+        self.token = token.access_token
+        self.expires_at = token.expires_at
+        self.expires_in = token.expires_in
+        self.refresh_expires_at = token.refresh_expires_at
+        self.refresh_expires_in = token.refresh_expires_in
+        self.refresh_token = token.refresh_token
+        self.token_type = token.token_type
+
+    def cache_credentials(self) -> None:
+        fpath = self._handle_creds_filepath(self.cache_dir)
+        with open(fpath, "w") as filebuffer:
+            json.dump(self._token.model_dump(), filebuffer, cls=TokenEncoder)
+
+    def refresh(self) -> None:
+        token: Token = self.client.refresh_token(self.token_type, self.token, self.refresh_token)
+        self._set_token(token)
+        self.cache_credentials()
+
+    def apply(self, headers: Mapping[str, str]) -> None:
+        """
+        Apply the Client secret key token to the X-CLIENT-SECRET header.
+        """
+        headers["Authorization"] = f"{self.token_type} {self.token}"

cmsdials/auth/client.py

@@ -0,0 +1,105 @@
+import time
+from datetime import datetime, timedelta
+
+import requests
+
+from ..utils.api_client import BaseAPIClient
+from ..utils.logger import logger
+from .exceptions import PendingAuthorizationError
+from .models import Device, DeviceToken, RefreshToken, Token
+
+
+class AuthClient(BaseAPIClient):
+    """
+    Client for DIALS auth endpoint
+    """
+
+    default_timeout = 30  # seconds
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+    def new_device(self) -> Device:
+        endpoint_url = self.api_url + "auth/new-device/"
+        response = requests.get(
+            url=endpoint_url,
+            headers={
+                "accept": "application/json",
+                "User-Agent": self._build_user_agent()
+            },
+            timeout=self.default_timeout,
+        )
+        response.raise_for_status()
+        response = response.json()
+        return Device(**response)
+
+    def device_token(self, device_code: str) -> Token:
+        endpoint_url = self.api_url + "auth/device-token/"
+        response = requests.post(
+            url=endpoint_url,
+            headers={
+                "accept": "application/json",
+                "Content-Type": "application/json",
+                "User-Agent": self._build_user_agent()
+            },
+            json={"device_code": device_code},
+            timeout=self.default_timeout,
+        )
+
+        if response.status_code == 400 and "authorization_pending" in response.text:
+            raise PendingAuthorizationError
+
+        response.raise_for_status()
+        response = response.json()
+        curr_datetime = datetime.now()
+        token = DeviceToken(**response)
+        return Token(
+            **token.dict(),
+            expires_at=curr_datetime + timedelta(seconds=token.expires_in),
+            refresh_expires_at=curr_datetime + timedelta(seconds=token.refresh_expires_in),
+        )
+
+    def refresh_token(self, token_type: str, access_token: str, refresh_token: str) -> Token:
+        endpoint_url = self.api_url + "auth/refresh-token/"
+        response = requests.post(
+            url=endpoint_url,
+            headers={
+                "accept": "application/json",
+                "Content-Type": "application/json",
+                "Authorization": f"{token_type} {access_token}",
+                "User-Agent": self._build_user_agent()
+            },
+            json={"refresh_token": refresh_token},
+            timeout=self.default_timeout,
+        )
+        response.raise_for_status()
+        response = response.json()
+        curr_datetime = datetime.now()
+        token = RefreshToken(**response)
+        return Token(
+            **token.dict(),
+            expires_at=curr_datetime + timedelta(seconds=token.expires_in),
+            refresh_expires_at=curr_datetime + timedelta(seconds=token.refresh_expires_in),
+        )
+
+    def device_auth_flow(self):
+        interval = 5  # seconds
+        device = self.new_device()
+        end_date = datetime.now() + timedelta(seconds=device.expires_in)
+
+        logger.info(f"This device will expire in {device.expires_in} seconds.")
+        logger.info(f"Go to the following url and authenticate: {device.verification_uri_complete}")
+        logger.info(f"Checking authorization status every {interval} seconds...")
+
+        while True:
+            time.sleep(interval)
+            curr_date = datetime.now()
+            if curr_date > end_date:
+                raise Exception("Action took too much time, device expired!")
+
+            try:
+                token = self.device_token(device_code=device.device_code)
+                logger.info("Device authorized, authentication finished successfully!")
+                return token
+            except PendingAuthorizationError:
+                logger.info("Device not authorized yet.")

cmsdials/auth/exceptions.py

@@ -0,0 +1,6 @@
+class PendingAuthorizationError(Exception):
+    pass
+
+
+class ImpossibleToRefreshTokenError(Exception):
+    pass

cmsdials/auth/models.py

@@ -0,0 +1,25 @@
+from datetime import datetime
+
+from pydantic import BaseModel
+
+
+class Device(BaseModel):
+    device_code: str
+    verification_uri_complete: str
+    expires_in: int
+
+
+class DeviceToken(BaseModel):
+    access_token: str
+    expires_in: int
+    refresh_expires_in: int
+    refresh_token: str
+    token_type: str
+
+
+RefreshToken = DeviceToken
+
+
+class Token(DeviceToken):
+    expires_at: datetime
+    refresh_expires_at: datetime

cmsdials/auth/secret_key.py

@@ -0,0 +1,42 @@
+from collections.abc import Mapping
+
+from ._base import BaseCredentials
+
+
+class Credentials(BaseCredentials):
+    """
+    Client secret key credentials.
+    These credential use Client secret key to provide authorization to applications.
+    """
+
+    def __init__(self, token: str) -> None:
+        if not token:
+            raise ValueError("token must be a non-empty Client secret key string")
+        self.token = token
+
+    @property
+    def stale(self) -> bool:
+        return False
+
+    @property
+    def expired(self) -> bool:
+        return False
+
+    @property
+    def refresh_expired(self) -> bool:
+        return False
+
+    def refresh(self) -> None:
+        return
+
+    def apply(self, headers: Mapping[str, str]) -> None:
+        """
+        Apply the Client secret key token to the X-CLIENT-SECRET header.
+        """
+        headers["X-CLIENT-SECRET"] = self.token
+
+    def before_request(self, headers: Mapping[str, str]):
+        """
+        Performs credential-specific before request logic.
+        """
+        self.apply(headers)

cmsdials/clients/dataset_index/client.py

@@ -0,0 +1,13 @@
+from ...utils.api_client import BaseAuthorizedAPIClient
+from .models import DatasetIndex, DatasetIndexFilters, PaginatedDatasetIndexList
+
+
+class DatasetIndexClient(BaseAuthorizedAPIClient):
+    data_model = DatasetIndex
+    pagination_model = PaginatedDatasetIndexList
+    filter_class = DatasetIndexFilters
+    lookup_url = "dataset-index/"
+
+    def get(self, dataset_id: int, file_id: int, **kwargs):
+        edp = f"{dataset_id}/{file_id}/"
+        return super().get(edp, **kwargs)

cmsdials/clients/dataset_index/models.py

@@ -0,0 +1,30 @@
+from datetime import datetime
+from typing import Optional
+
+from pydantic import AnyUrl, BaseModel, Field
+
+from ...utils.base_model import OBaseModel, PaginatedBaseModel
+
+
+class DatasetIndex(BaseModel):
+    dataset_id: int
+    dataset: str = Field(..., max_length=255)
+    era: str = Field(..., max_length=255)
+    data_tier: str = Field(..., max_length=255)
+    primary_ds_name: str = Field(..., max_length=255)
+    processed_ds_name: str = Field(..., max_length=255)
+    processing_version: int
+    last_modification_date: datetime
+
+
+class PaginatedDatasetIndexList(PaginatedBaseModel):
+    next: Optional[AnyUrl]
+    previous: Optional[AnyUrl]
+    results: list[DatasetIndex]
+
+
+class DatasetIndexFilters(OBaseModel):
+    next_token: Optional[str] = None
+    page_size: Optional[int] = None
+    dataset: Optional[str] = None
+    dataset__regex: Optional[str] = None

cmsdials/clients/file_index/client.py

@@ -0,0 +1,13 @@
+from ...utils.api_client import BaseAuthorizedAPIClient
+from .models import FileIndex, FileIndexFilters, PaginatedFileIndexList
+
+
+class FileIndexClient(BaseAuthorizedAPIClient):
+    data_model = FileIndex
+    pagination_model = PaginatedFileIndexList
+    filter_class = FileIndexFilters
+    lookup_url = "file-index/"
+
+    def get(self, dataset_id: int, file_id: int, **kwargs):
+        edp = f"{dataset_id}/{file_id}/"
+        return super().get(edp, **kwargs)

cmsdials/clients/file_index/models.py

@@ -0,0 +1,36 @@
+from datetime import datetime
+from typing import Optional
+
+from pydantic import AnyUrl, BaseModel, Field
+
+from ...utils.base_model import OBaseModel, PaginatedBaseModel
+
+
+class FileIndex(BaseModel):
+    dataset_id: int
+    dataset: str = Field(..., max_length=255)
+    file_id: int
+    file_size: int
+    creation_date: datetime
+    last_modification_date: datetime
+    logical_file_name: str
+    status: str = Field(..., max_length=255)
+    err_trace: Optional[str] = Field(..., max_length=5000)
+
+
+class PaginatedFileIndexList(PaginatedBaseModel):
+    next: Optional[AnyUrl]
+    previous: Optional[AnyUrl]
+    results: list[FileIndex]
+
+
+class FileIndexFilters(OBaseModel):
+    next_token: Optional[str] = None
+    page_size: Optional[int] = None
+    dataset_id: Optional[int] = None
+    logical_file_name: Optional[str] = None
+    logical_file_name__regex: Optional[str] = None
+    status: Optional[str] = None
+    min_size: Optional[int] = None
+    dataset: Optional[str] = None
+    dataset__regex: Optional[str] = None

cmsdials/clients/h1d/client.py

@@ -0,0 +1,13 @@
+from ...utils.api_client import BaseAuthorizedAPIClient
+from .models import LumisectionHistogram1D, LumisectionHistogram1DFilters, PaginatedLumisectionHistogram1DList
+
+
+class LumisectionHistogram1DClient(BaseAuthorizedAPIClient):
+    data_model = LumisectionHistogram1D
+    pagination_model = PaginatedLumisectionHistogram1DList
+    filter_class = LumisectionHistogram1DFilters
+    lookup_url = "th1/"
+
+    def get(self, dataset_id: int, run_number: int, ls_number: int, me_id: int, **kwargs):
+        edp = f"{dataset_id}/{run_number}/{ls_number}/{me_id}/"
+        return super().get(edp, **kwargs)

cmsdials/clients/h1d/models.py

@@ -0,0 +1,47 @@
+from typing import Optional
+
+from pydantic import AnyUrl, BaseModel, Field
+
+from ...utils.base_model import OBaseModel, PaginatedBaseModel
+
+
+class LumisectionHistogram1D(BaseModel):
+    dataset: str = Field(..., max_length=255)
+    me: str = Field(..., max_length=255)
+    dataset_id: int
+    file_id: int
+    run_number: int
+    ls_number: int
+    me_id: int
+    x_min: float
+    x_max: float
+    x_bin: float
+    entries: int
+    data: list[float]
+
+
+class PaginatedLumisectionHistogram1DList(PaginatedBaseModel):
+    next: Optional[AnyUrl]
+    previous: Optional[AnyUrl]
+    results: list[LumisectionHistogram1D]
+
+
+class LumisectionHistogram1DFilters(OBaseModel):
+    next_token: Optional[str] = None
+    page_size: Optional[int] = None
+    dataset_id: Optional[int] = None
+    file_id: Optional[int] = None
+    run_number: Optional[int] = None
+    run_number__lte: Optional[int] = None
+    run_number__gte: Optional[int] = None
+    ls_number: Optional[int] = None
+    ls_numbet__lte: Optional[int] = None
+    ls_number__gte: Optional[int] = None
+    me_id: Optional[int] = None
+    entries__gte: Optional[int] = None
+    dataset: Optional[str] = None
+    dataset__regex: Optional[str] = None
+    logical_file_name: Optional[str] = None
+    logical_file_name__regex: Optional[str] = None
+    me: Optional[str] = None
+    me__regex: Optional[str] = None

cmsdials/clients/h2d/client.py

@@ -0,0 +1,13 @@
+from ...utils.api_client import BaseAuthorizedAPIClient
+from .models import LumisectionHistogram2D, LumisectionHistogram2DFilters, PaginatedLumisectionHistogram2DList
+
+
+class LumisectionHistogram2DClient(BaseAuthorizedAPIClient):
+    data_model = LumisectionHistogram2D
+    pagination_model = PaginatedLumisectionHistogram2DList
+    filter_class = LumisectionHistogram2DFilters
+    lookup_url = "th2/"
+
+    def get(self, dataset_id: int, run_number: int, ls_number: int, me_id: int, **kwargs):
+        edp = f"{dataset_id}/{run_number}/{ls_number}/{me_id}/"
+        return super().get(edp, **kwargs)