rucio 32.8.6__py3-none-any.whl → 35.8.0__py3-none-any.whl

rucio/core/account_limit.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # Copyright European Organization for Nuclear Research (CERN) since 2012
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,8 +14,8 @@
 
 from typing import TYPE_CHECKING
 
-from sqlalchemy.orm.exc import NoResultFound
-from sqlalchemy.sql import func, select, literal
+from sqlalchemy.exc import NoResultFound
+from sqlalchemy.sql import func, literal, select
 from sqlalchemy.sql.expression import and_, or_
 
 from rucio.core.account import get_all_rse_usages_per_account
@@ -99,9 +98,17 @@ def get_global_account_limits(account=None, *, session: "Session"):
     :return:         Dict {'MOCK': {'resolved_rses': ['MOCK'], 'limit': 10, 'resolved_rse_ids': [123]}}.
     """
     if account:
-        global_account_limits = session.query(models.AccountGlobalLimit).filter_by(account=account).all()
+        stmt = select(
+            models.AccountGlobalLimit
+        ).where(
+            models.AccountGlobalLimit.account == account
+        )
+        global_account_limits = session.execute(stmt).scalars().all()
     else:
-        global_account_limits = session.query(models.AccountGlobalLimit).all()
+        stmt = select(
+            models.AccountGlobalLimit
+        )
+        global_account_limits = session.execute(stmt).scalars().all()
 
     resolved_global_account_limits = {}
     for limit in global_account_limits:
@@ -131,7 +138,13 @@ def get_global_account_limit(account, rse_expression, *, session: "Session"):
     :return:                Limit in Bytes.
     """
     try:
-        global_account_limit = session.query(models.AccountGlobalLimit).filter_by(account=account, rse_expression=rse_expression).one()
+        stmt = select(
+            models.AccountGlobalLimit
+        ).where(
+            and_(models.AccountGlobalLimit.account == account,
+                 models.AccountGlobalLimit.rse_expression == rse_expression)
+        )
+        global_account_limit = session.execute(stmt).scalar_one()
         if global_account_limit.bytes == -1:
             return float("inf")
         else:
@@ -151,8 +164,13 @@ def get_local_account_limit(account, rse_id, *, session: "Session"):
     :return:         Limit in Bytes.
     """
     try:
-        account_limit = session.query(models.AccountLimit).filter(models.AccountLimit.account == account,
-                                                                  models.AccountLimit.rse_id == rse_id).one()
+        stmt = select(
+            models.AccountLimit
+        ).where(
+            and_(models.AccountLimit.account == account,
+                 models.AccountLimit.rse_id == rse_id)
+        )
+        account_limit = session.execute(stmt).scalar_one()
         if account_limit.bytes == -1:
             return float("inf")
         else:
@@ -176,17 +194,28 @@ def get_local_account_limits(account, rse_ids=None, *, session: "Session"):
     if rse_ids:
         rse_id_clauses = []
         for rse_id in rse_ids:
-            rse_id_clauses.append(and_(models.AccountLimit.rse_id == rse_id, models.AccountLimit.account == account))
+            rse_id_clauses.append(and_(models.AccountLimit.rse_id == rse_id,
+                                       models.AccountLimit.account == account))
         rse_id_clause_chunks = [rse_id_clauses[x:x + 10] for x in range(0, len(rse_id_clauses), 10)]
         for rse_id_chunk in rse_id_clause_chunks:
-            tmp_limits = session.query(models.AccountLimit).filter(or_(*rse_id_chunk)).all()
+            stmt = select(
+                models.AccountLimit
+            ).where(
+                or_(*rse_id_chunk)
+            )
+            tmp_limits = session.execute(stmt).scalars().all()
             for limit in tmp_limits:
                 if limit.bytes == -1:
                     account_limits[limit.rse_id] = float("inf")
                 else:
                     account_limits[limit.rse_id] = limit.bytes
     else:
-        account_limits_tmp = session.query(models.AccountLimit).filter(models.AccountLimit.account == account).all()
+        stmt = select(
+            models.AccountLimit
+        ).where(
+            models.AccountLimit.account == account
+        )
+        account_limits_tmp = session.execute(stmt).scalars().all()
         for limit in account_limits_tmp:
             if limit.bytes == -1:
                 account_limits[limit.rse_id] = float("inf")
@@ -206,8 +235,13 @@ def set_local_account_limit(account, rse_id, bytes_, *, session: "Session"):
     :param session:  Database session in use.
     """
     try:
-        account_limit = session.query(models.AccountLimit).filter(models.AccountLimit.account == account,
-                                                                  models.AccountLimit.rse_id == rse_id).one()
+        stmt = select(
+            models.AccountLimit
+        ).where(
+            and_(models.AccountLimit.account == account,
+                 models.AccountLimit.rse_id == rse_id)
+        )
+        account_limit = session.execute(stmt).scalar_one()
         account_limit.bytes = bytes_
     except NoResultFound:
         models.AccountLimit(account=account, rse_id=rse_id, bytes=bytes_).save(session=session)
@@ -224,8 +258,13 @@ def set_global_account_limit(account, rse_expression, bytes_, *, session: "Sessi
     :param session:         Database session in use.
     """
     try:
-        account_limit = session.query(models.AccountGlobalLimit).filter(models.AccountGlobalLimit.account == account,
-                                                                        models.AccountGlobalLimit.rse_expression == rse_expression).one()
+        stmt = select(
+            models.AccountGlobalLimit
+        ).where(
+            and_(models.AccountGlobalLimit.account == account,
+                 models.AccountGlobalLimit.rse_expression == rse_expression)
+        )
+        account_limit = session.execute(stmt).scalar_one()
         account_limit.bytes = bytes_
     except NoResultFound:
         models.AccountGlobalLimit(account=account, rse_expression=rse_expression, bytes=bytes_).save(session=session)
@@ -242,8 +281,14 @@ def delete_local_account_limit(account, rse_id, *, session: "Session"):
     :returns:        True if something was deleted; False otherwise.
     """
     try:
-        session.query(models.AccountLimit).filter(models.AccountLimit.account == account,
-                                                  models.AccountLimit.rse_id == rse_id).one().delete(session=session)
+        stmt = select(
+            models.AccountLimit
+        ).where(
+            and_(models.AccountLimit.account == account,
+                 models.AccountLimit.rse_id == rse_id)
+        )
+        result = session.execute(stmt).scalar_one()
+        result.delete(session=session)
         return True
     except NoResultFound:
         return False
@@ -260,8 +305,14 @@ def delete_global_account_limit(account, rse_expression, *, session: "Session"):
     :returns:               True if something was deleted; False otherwise.
     """
     try:
-        session.query(models.AccountGlobalLimit).filter(models.AccountGlobalLimit.account == account,
-                                                        models.AccountGlobalLimit.rse_expression == rse_expression).one().delete(session=session)
+        stmt = select(
+            models.AccountGlobalLimit
+        ).where(
+            and_(models.AccountGlobalLimit.account == account,
+                 models.AccountGlobalLimit.rse_expression == rse_expression)
+        )
+        result = session.execute(stmt).scalar_one()
+        result.delete(session=session)
         return True
     except NoResultFound:
         return False
@@ -279,14 +330,22 @@ def get_local_account_usage(account, rse_id=None, *, session: "Session"):
     :returns:        List of dicts {'rse_id', 'rse', 'bytes', 'files', 'bytes_limit', 'bytes_remaining'}
     """
 
+    stmt = select(
+        models.AccountUsage
+    ).where(
+        models.AccountUsage.account == account
+    )
     if not rse_id:
         # All RSESs
         limits = get_local_account_limits(account=account, session=session)
-        counters = {c.rse_id: c for c in session.query(models.AccountUsage).filter_by(account=account).all()}
+        counters = {c.rse_id: c for c in session.execute(stmt).scalars().all()}
     else:
         # One RSE
+        stmt.where(
+            models.AccountUsage.rse_id == rse_id
+        )
         limits = get_local_account_limits(account=account, rse_ids=[rse_id], session=session)
-        counters = {c.rse_id: c for c in session.query(models.AccountUsage).filter_by(account=account, rse_id=rse_id).all()}
+        counters = {c.rse_id: c for c in session.execute(stmt).scalars().all()}
     result_list = []
 
     for rse_id in set(limits).union(counters):
@@ -341,9 +400,16 @@ def get_global_account_usage(account, rse_expression=None, *, session: "Session"
         limit = get_global_account_limit(account=account, rse_expression=rse_expression, session=session)
         vo = account.vo
         resolved_rses = [resolved_rse['id'] for resolved_rse in parse_expression(rse_expression, filter_={'vo': vo}, session=session)]
-        usage = session.query(func.sum(models.AccountUsage.bytes), func.sum(models.AccountUsage.files))\
-                       .filter(models.AccountUsage.account == account, models.AccountUsage.rse_id.in_(resolved_rses))\
-                       .group_by(models.AccountUsage.account).first()
+        stmt = select(
+            func.sum(models.AccountUsage.bytes),
+            func.sum(models.AccountUsage.files)
+        ).where(
+            and_(models.AccountUsage.account == account,
+                 models.AccountUsage.rse_id.in_(resolved_rses))
+        ).group_by(
+            models.AccountUsage.account
+        )
+        usage = session.execute(stmt).first()
         if limit is None:
             limit = 0
         if usage is None:

rucio/core/authentication.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # Copyright European Organization for Nuclear Research (CERN) since 2012
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,27 +19,27 @@ import re
 import sys
 import traceback
 from base64 import b64decode
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional, Union, cast
 
 import paramiko
 from dogpile.cache import make_region
-from dogpile.cache.api import NO_VALUE
+from dogpile.cache.api import NO_VALUE, NoValue
 from sqlalchemy import delete, null, or_, select
 
 from rucio.common.cache import make_region_memcached
 from rucio.common.config import config_get_bool
 from rucio.common.exception import CannotAuthenticate, RucioException
-from rucio.common.utils import chunks, generate_uuid, date_to_str
+from rucio.common.utils import chunks, date_to_str, generate_uuid
 from rucio.core.account import account_exists
 from rucio.core.oidc import validate_jwt
-from rucio.db.sqla import filter_thread_work
-from rucio.db.sqla import models
+from rucio.db.sqla import filter_thread_work, models
 from rucio.db.sqla.constants import IdentityType
 from rucio.db.sqla.session import read_session, transactional_session
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
-    from typing import Any, Union
+
+    from rucio.common.types import InternalAccount, TokenDict, TokenValidationDict
 
 
 def strip_x509_proxy_attributes(dn: str) -> str:
@@ -90,7 +89,15 @@ else:
 
 
 @transactional_session
-def get_auth_token_user_pass(account, username, password, appid, ip=None, *, session: "Session"):
+def get_auth_token_user_pass(
+    account: "InternalAccount",
+    username: str,
+    password: str,
+    appid: str,
+    ip: Optional[str] = None,
+    *,
+    session: "Session"
+) -> Optional["TokenDict"]:
     """
     Authenticate a Rucio account temporarily via username and password.
 
@@ -150,7 +157,14 @@ def get_auth_token_user_pass(account, username, password, appid, ip=None, *, ses
 
 
 @transactional_session
-def get_auth_token_x509(account, dn, appid, ip=None, *, session: "Session"):
+def get_auth_token_x509(
+    account: "InternalAccount",
+    dn: str,
+    appid: str,
+    ip: Optional[str] = None,
+    *,
+    session: "Session"
+) -> Optional["TokenDict"]:
     """
     Authenticate a Rucio account temporarily via an x509 certificate.
 
@@ -182,7 +196,14 @@ def get_auth_token_x509(account, dn, appid, ip=None, *, session: "Session"):
 
 
 @transactional_session
-def get_auth_token_gss(account, gsstoken, appid, ip=None, *, session: "Session"):
+def get_auth_token_gss(
+    account: "InternalAccount",
+    gsstoken: str,
+    appid: str,
+    ip: Optional[str] = None,
+    *,
+    session: "Session"
+) -> Optional["TokenDict"]:
     """
     Authenticate a Rucio account temporarily via a GSS token.
 
@@ -214,22 +235,34 @@ def get_auth_token_gss(account, gsstoken, appid, ip=None, *, session: "Session")
 
 
 @transactional_session
-def get_auth_token_ssh(account, signature, appid, ip=None, *, session: "Session"):
+def get_auth_token_ssh(
+    account: "InternalAccount",
+    signature: str,
+    appid: str,
+    ip: Optional[str] = None,
+    *,
+    session: "Session"
+) -> Optional["TokenDict"]:
     """
     Authenticate a Rucio account temporarily via SSH key exchange.
 
     The token lifetime is 1 hour.
 
     :param account: Account identifier as a string.
-    :param signature: Response to server challenge signed with SSH private key as string.
+    :param signature: Response to server challenge signed with SSH private key as a base64 encoded string.
     :param appid: The application identifier as a string.
     :param ip: IP address of the client as a string.
     :param session: The database session in use.
 
     :returns: A dict with token and expires_at entries.
     """
-    if not isinstance(signature, bytes):
-        signature = signature.encode()
+
+    # decode the signature which must come in base64 encoded
+    try:
+        signature += '=' * ((4 - len(signature) % 4) % 4)  # adding required padding
+        decoded_signature = b64decode(signature)
+    except TypeError:
+        raise CannotAuthenticate(f'Cannot authenticate to account {account} with malformed signature')
 
     # Make sure the account exists
     if not account_exists(account, session=session):
@@ -266,7 +299,7 @@ def get_auth_token_ssh(account, signature, appid, ip=None, *, session: "Session"
         pub_k = paramiko.RSAKey(data=b64decode(data))
         for challenge_token in active_challenge_tokens:
             if pub_k.verify_ssh_sig(str(challenge_token['token']).encode(),
-                                    paramiko.Message(signature)):
+                                    paramiko.Message(decoded_signature)):
                 match = True
                 break
         if match:
@@ -288,7 +321,13 @@ def get_auth_token_ssh(account, signature, appid, ip=None, *, session: "Session"
 
 
 @transactional_session
-def get_ssh_challenge_token(account, appid, ip=None, *, session: "Session"):
+def get_ssh_challenge_token(
+    account: "InternalAccount",
+    appid: str,
+    ip: Optional[str] = None,
+    *,
+    session: "Session"
+) -> Optional["TokenDict"]:
     """
     Prepare a challenge token for subsequent SSH public key authentication.
 
@@ -324,7 +363,14 @@ def get_ssh_challenge_token(account, appid, ip=None, *, session: "Session"):
 
 
 @transactional_session
-def get_auth_token_saml(account, saml_nameid, appid, ip=None, *, session: "Session"):
+def get_auth_token_saml(
+    account: "InternalAccount",
+    saml_nameid: str,
+    appid: str,
+    ip: Optional[str] = None,
+    *,
+    session: "Session"
+) -> Optional["TokenDict"]:
     """
     Authenticate a Rucio account temporarily via SAML.
 
@@ -355,7 +401,12 @@ def get_auth_token_saml(account, saml_nameid, appid, ip=None, *, session: "Sessi
 
 
 @transactional_session
-def redirect_auth_oidc(auth_code, fetchtoken=False, *, session: "Session"):
+def redirect_auth_oidc(
+    auth_code: str,
+    fetchtoken: bool = False,
+    *,
+    session: "Session"
+) -> Optional[str]:
     """
     Finds the Authentication URL in the Rucio DB oauth_requests table
     and redirects user's browser to this URL.
@@ -396,7 +447,13 @@ def redirect_auth_oidc(auth_code, fetchtoken=False, *, session: "Session"):
 
 
 @transactional_session
-def delete_expired_tokens(total_workers, worker_number, limit=1000, *, session: "Session"):
+def delete_expired_tokens(
+    total_workers: int,
+    worker_number: int,
+    limit: int = 1000,
+    *,
+    session: "Session"
+) -> int:
     """
     Delete expired tokens.
 
@@ -456,7 +513,7 @@ def delete_expired_tokens(total_workers, worker_number, limit=1000, *, session:
 
 
 @read_session
-def query_token(token, *, session: "Session"):
+def query_token(token: str, *, session: "Session") -> Optional["TokenValidationDict"]:
     """
     Validate an authentication token using the database. This method will only be called
     if no entry could be found in the according cache.
@@ -484,12 +541,12 @@ def query_token(token, *, session: "Session"):
     )
     result = session.execute(query).first()
     if result:
-        return result._asdict()
+        return cast("TokenValidationDict", result._asdict())
     return None
 
 
 @transactional_session
-def validate_auth_token(token: str, *, session: "Session") -> "dict[str, Any]":
+def validate_auth_token(token: str, *, session: "Session") -> "TokenValidationDict":
     """
     Validate an authentication token.
 
@@ -510,12 +567,12 @@ def validate_auth_token(token: str, *, session: "Session") -> "dict[str, Any]":
     token = token.strip()
     cache_key = token.replace(' ', '')
 
-    # Check if token ca be found in cache region
-    value: "Union[NO_VALUE, dict[str, Any]]" = TOKENREGION.get(cache_key)
+    # Check if token can be found in cache region
+    value: Union[NoValue, "TokenValidationDict"] = TOKENREGION.get(cache_key)
     if value is NO_VALUE:  # no cached entry found
         value = query_token(token, session=session)
         if not value:
-            # identify JWT access token and validte
+            # identify JWT access token and validate
             # & save it in Rucio if scope and audience are correct
             if len(token.split(".")) == 3:
                 value = validate_jwt(token, session=session)
@@ -523,19 +580,19 @@ def validate_auth_token(token: str, *, session: "Session") -> "dict[str, Any]":
                 raise CannotAuthenticate(traceback.format_exc())
         # save token in the cache
         TOKENREGION.set(cache_key, value)
-    lifetime = value.get('lifetime', datetime.datetime(1970, 1, 1))
+    lifetime = value.get('lifetime', datetime.datetime(1970, 1, 1))  # type: ignore (value is narrowed to dict, but type-checker doesn't see it)
     if lifetime < datetime.datetime.utcnow():  # check if expired
         TOKENREGION.delete(cache_key)
         raise CannotAuthenticate(f"Token found but expired since {date_to_str(lifetime)}.")
-    return value
+    return cast("TokenValidationDict", value)
 
 
-def token_dictionary(token: models.Token):
+def token_dictionary(token: models.Token) -> "TokenDict":
     return {'token': token.token, 'expires_at': token.expired_at}
 
 
 @transactional_session
-def __delete_expired_tokens_account(account, *, session: "Session"):
+def __delete_expired_tokens_account(account: "InternalAccount", *, session: "Session") -> None:
     """"
     Deletes expired tokens from the database.