From fb1bc8acf50d7f91371c1d50c84c40430cb9f006 Mon Sep 17 00:00:00 2001 From: Chris Burr Date: Sat, 27 Jan 2024 16:00:04 +0100 Subject: [PATCH] feat: Add ProxyDB --- diracx-core/src/diracx/core/exceptions.py | 8 ++ diracx-db/pyproject.toml | 1 + diracx-db/src/diracx/db/sql/__init__.py | 10 +- diracx-db/src/diracx/db/sql/proxy/__init__.py | 0 diracx-db/src/diracx/db/sql/proxy/db.py | 120 ++++++++++++++++ diracx-db/src/diracx/db/sql/proxy/schema.py | 21 +++ diracx-db/tests/proxy/test_proxydb.py | 129 ++++++++++++++++++ .../src/diracx/routers/dependencies.py | 3 + 8 files changed, 291 insertions(+), 1 deletion(-) create mode 100644 diracx-db/src/diracx/db/sql/proxy/__init__.py create mode 100644 diracx-db/src/diracx/db/sql/proxy/db.py create mode 100644 diracx-db/src/diracx/db/sql/proxy/schema.py create mode 100644 diracx-db/tests/proxy/test_proxydb.py diff --git a/diracx-core/src/diracx/core/exceptions.py b/diracx-core/src/diracx/core/exceptions.py index 65d143e07..19561382b 100644 --- a/diracx-core/src/diracx/core/exceptions.py +++ b/diracx-core/src/diracx/core/exceptions.py @@ -43,3 +43,11 @@ class JobNotFound(Exception): def __init__(self, job_id: int): self.job_id: int = job_id super().__init__(f"Job {job_id} not found") + + +class VOMSInitError(DiracError): + """Adding VOMS attributes to a proxy failed""" + + +class ProxyNotFoundError(DiracError): + """There are no valid proxies for the given user""" diff --git a/diracx-db/pyproject.toml b/diracx-db/pyproject.toml index d400b3e23..380247c4d 100644 --- a/diracx-db/pyproject.toml +++ b/diracx-db/pyproject.toml @@ -31,6 +31,7 @@ testing = [ AuthDB = "diracx.db.sql:AuthDB" JobDB = "diracx.db.sql:JobDB" JobLoggingDB = "diracx.db.sql:JobLoggingDB" +ProxyDB = "diracx.db.sql:ProxyDB" SandboxMetadataDB = "diracx.db.sql:SandboxMetadataDB" TaskQueueDB = "diracx.db.sql:TaskQueueDB" diff --git a/diracx-db/src/diracx/db/sql/__init__.py b/diracx-db/src/diracx/db/sql/__init__.py index 582509b13..a3a2d588c 100644 --- a/diracx-db/src/diracx/db/sql/__init__.py +++ b/diracx-db/src/diracx/db/sql/__init__.py @@ -1,7 +1,15 @@ from __future__ import annotations -__all__ = ("AuthDB", "JobDB", "JobLoggingDB", "SandboxMetadataDB", "TaskQueueDB") +__all__ = ( + "AuthDB", + "JobDB", + "JobLoggingDB", + "ProxyDB", + "SandboxMetadataDB", + "TaskQueueDB", +) from .auth.db import AuthDB from .jobs.db import JobDB, JobLoggingDB, TaskQueueDB +from .proxy.db import ProxyDB from .sandbox_metadata.db import SandboxMetadataDB diff --git a/diracx-db/src/diracx/db/sql/proxy/__init__.py b/diracx-db/src/diracx/db/sql/proxy/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/diracx-db/src/diracx/db/sql/proxy/db.py b/diracx-db/src/diracx/db/sql/proxy/db.py new file mode 100644 index 000000000..77a602227 --- /dev/null +++ b/diracx-db/src/diracx/db/sql/proxy/db.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import asyncio +import os +import stat +from datetime import datetime, timezone +from pathlib import Path +from subprocess import DEVNULL, PIPE, STDOUT +from tempfile import TemporaryDirectory + +from DIRAC.Core.Security import Locations +from DIRAC.Core.Security.VOMS import voms_init_cmd +from DIRAC.Core.Security.X509Chain import X509Chain +from DIRAC.Core.Utilities.ReturnValues import returnValueOrRaise +from sqlalchemy import select + +from diracx.core.exceptions import ProxyNotFoundError, VOMSInitError +from diracx.db.sql.utils import BaseSQLDB, utcnow + +from .schema import Base as ProxyDBBase +from .schema import CleanProxies + +PROXY_PROVIDER = "Certificate" + + +class ProxyDB(BaseSQLDB): + metadata = ProxyDBBase.metadata + + async def get_proxy( + self, + dn: str, + vo: str, + dirac_group: str, + voms_attr: str | None, + lifetime_seconds: int, + ) -> str: + """Generate a new proxy for the given DN as PEM with the given VOMS extension""" + original_chain = await self.get_stored_proxy( + dn, min_lifetime_seconds=lifetime_seconds + ) + + proxy_string = returnValueOrRaise( + original_chain.generateProxyToString( + lifetime_seconds, + diracGroup=dirac_group, + strength=returnValueOrRaise(original_chain.getStrength()), + ) + ) + proxy_chain = X509Chain() + proxy_chain.loadProxyFromString(proxy_string) + + with TemporaryDirectory() as tmpdir: + in_fn = Path(tmpdir) / "in.pem" + in_fn.touch(stat.S_IRUSR | stat.S_IWUSR) + in_fn.write_text(proxy_string) + out_fn = Path(tmpdir) / "out.pem" + + cmd = voms_init_cmd( + vo, + voms_attr, + proxy_chain, + str(in_fn), + str(out_fn), + Locations.getVomsesLocation(), + ) + proc = await asyncio.create_subprocess_exec( + *cmd, + stdin=DEVNULL, + stdout=PIPE, + stderr=STDOUT, + env=os.environ + | { + "X509_CERT_DIR": Locations.getCAsLocationNoConfig(), + "X509_VOMS_DIR": Locations.getVomsdirLocation(), + }, + ) + await proc.wait() + if proc.returncode != 0: + assert proc.stdout + message = (await proc.stdout.read()).decode("utf-8", "backslashreplace") + raise VOMSInitError( + f"voms-proxy-init failed with return code {proc.returncode}: {message}" + ) + + voms_string = out_fn.read_text() + + return voms_string + + async def get_stored_proxy( + self, dn: str, *, min_lifetime_seconds: int + ) -> X509Chain: + """Get the X509 proxy that is stored in the DB for the given DN + + NOTE: This is the original long-lived proxy and should only be used to + generate short-lived proxies!!! + """ + stmt = select(CleanProxies.Pem, CleanProxies.ExpirationTime) + stmt = stmt.where( + CleanProxies.UserDN == dn, + CleanProxies.ExpirationTime > utcnow(), + CleanProxies.ProxyProvider == PROXY_PROVIDER, + ) + + for pem_data, expiration_time in (await self.conn.execute(stmt)).all(): + seconds_remaining = ( + expiration_time.replace(tzinfo=timezone.utc) + - datetime.now(timezone.utc) + ).total_seconds() + if seconds_remaining <= min_lifetime_seconds: + continue + + pem_data = pem_data.decode("ascii") + if not pem_data: + continue + chain = X509Chain() + returnValueOrRaise(chain.loadProxyFromString(pem_data)) + return chain + raise ProxyNotFoundError( + f"No proxy found for {dn} with over {min_lifetime_seconds} seconds of life" + ) diff --git a/diracx-db/src/diracx/db/sql/proxy/schema.py b/diracx-db/src/diracx/db/sql/proxy/schema.py new file mode 100644 index 000000000..23cf2646c --- /dev/null +++ b/diracx-db/src/diracx/db/sql/proxy/schema.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +from sqlalchemy import ( + BLOB, + DateTime, + String, +) +from sqlalchemy.orm import declarative_base + +from diracx.db.sql.utils import Column, NullColumn + +Base = declarative_base() + + +class CleanProxies(Base): + __tablename__ = "ProxyDB_CleanProxies" + UserName = Column(String(64)) + Pem = NullColumn(BLOB) + ProxyProvider = Column(String(64), default="Certificate") + ExpirationTime = NullColumn(DateTime) + UserDN = Column(String(255), primary_key=True) diff --git a/diracx-db/tests/proxy/test_proxydb.py b/diracx-db/tests/proxy/test_proxydb.py new file mode 100644 index 000000000..a59830fe2 --- /dev/null +++ b/diracx-db/tests/proxy/test_proxydb.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from functools import wraps +from pathlib import Path +from typing import AsyncGenerator + +import pytest +from DIRAC.Core.Security.VOMS import voms_init_cmd +from DIRAC.Core.Security.X509Chain import X509Chain +from DIRAC.Core.Utilities.ReturnValues import returnValueOrRaise +from sqlalchemy import insert + +from diracx.core.exceptions import DiracError +from diracx.db.sql.proxy.db import ProxyDB +from diracx.db.sql.proxy.schema import CleanProxies + +TEST_NAME = "testuser" +TEST_DN = "/O=Dirac Computing/O=CERN/CN=MrUser" +TEST_DATA_DIR = Path(__file__).parent / "data" +TEST_PEM_PATH = TEST_DATA_DIR / "proxy.pem" + + +@pytest.fixture +async def empty_proxy_db(tmp_path) -> AsyncGenerator[ProxyDB, None]: + proxy_db = ProxyDB("sqlite+aiosqlite:///:memory:") + async with proxy_db.engine_context(): + async with proxy_db.engine.begin() as conn: + await conn.run_sync(proxy_db.metadata.create_all) + yield proxy_db + + +@pytest.fixture +async def proxy_db(empty_proxy_db) -> AsyncGenerator[ProxyDB, None]: + async with empty_proxy_db.engine.begin() as conn: + await conn.execute( + insert(CleanProxies).values( + UserName=TEST_NAME, + UserDN=TEST_DN, + ProxyProvider="Certificate", + Pem=TEST_PEM_PATH.read_bytes(), + ExpirationTime=datetime(2033, 11, 25, 21, 25, 23, tzinfo=timezone.utc), + ) + ) + yield empty_proxy_db + + +async def test_get_stored_proxy(proxy_db: ProxyDB): + async with proxy_db as proxy_db: + proxy = await proxy_db.get_stored_proxy(TEST_DN, min_lifetime_seconds=3600) + assert proxy + + +async def test_no_proxy_for_dn_1(empty_proxy_db: ProxyDB): + async with empty_proxy_db as proxy_db: + with pytest.raises(DiracError, match="No proxy found"): + await proxy_db.get_stored_proxy(TEST_DN, min_lifetime_seconds=3600) + + +async def test_no_proxy_for_dn_2(empty_proxy_db: ProxyDB): + async with empty_proxy_db as proxy_db: + with pytest.raises(DiracError, match="No proxy found"): + await proxy_db.get_stored_proxy( + "/O=OtherOrg/O=CERN/CN=MrUser", min_lifetime_seconds=3600 + ) + + +async def test_proxy_not_long_enough(proxy_db: ProxyDB): + async with proxy_db as proxy_db: + with pytest.raises(DiracError, match="No proxy found"): + # The test proxy we use is valid for 10 years + # If this code still exists in 2028 we might start having problems with 2K38 + await proxy_db.get_stored_proxy( + TEST_DN, min_lifetime_seconds=10 * 365 * 24 * 3600 + ) + + +@wraps(voms_init_cmd) +def voms_init_cmd_fake(*args, **kwargs): + cmd = voms_init_cmd(*args, **kwargs) + + new_cmd = ["voms-proxy-fake"] + i = 1 + while i < len(cmd): + # Some options are not supported by voms-proxy-fake + if cmd[i] in {"-valid", "-vomses", "-timeout"}: + i += 2 + continue + new_cmd.append(cmd[i]) + i += 1 + new_cmd.extend( + [ + "-hostcert", + f"{TEST_DATA_DIR}/certs/host/hostcert.pem", + "-hostkey", + f"{TEST_DATA_DIR}/certs/host/hostkey.pem", + "-fqan", + "/fakevo/Role=NULL/Capability=NULL", + ] + ) + return new_cmd + + +async def test_get_proxy(proxy_db: ProxyDB, monkeypatch): + monkeypatch.setenv("X509_CERT_DIR", str(TEST_DATA_DIR / "certs")) + monkeypatch.setattr("diracx.db.sql.proxy.db.voms_init_cmd", voms_init_cmd_fake) + + async with proxy_db as proxy_db: + proxy_pem = await proxy_db.get_proxy( + TEST_DN, "fakevo", "fakevo_user", "/fakevo", 3600 + ) + + proxy_chain = X509Chain() + returnValueOrRaise(proxy_chain.loadProxyFromString(proxy_pem)) + + # Check validity + not_after = returnValueOrRaise(proxy_chain.getNotAfterDate()).replace( + tzinfo=timezone.utc + ) + # The proxy should currently be valid + assert datetime.now(timezone.utc) < not_after + # The proxy should be invalid in less than 3601 seconds + time_left = not_after - datetime.now(timezone.utc) + assert time_left < timedelta(hours=1, seconds=1) + + # Check VOMS data + voms_data = returnValueOrRaise(proxy_chain.getVOMSData()) + assert voms_data["vo"] == "fakevo" + assert voms_data["fqan"] == ["/fakevo/Role=NULL/Capability=NULL"] diff --git a/diracx-routers/src/diracx/routers/dependencies.py b/diracx-routers/src/diracx/routers/dependencies.py index 7a67b94f4..c3bdd18a2 100644 --- a/diracx-routers/src/diracx/routers/dependencies.py +++ b/diracx-routers/src/diracx/routers/dependencies.py @@ -5,6 +5,7 @@ "AuthDB", "JobDB", "JobLoggingDB", + "ProxyDB", "SandboxMetadataDB", "TaskQueueDB", "add_settings_annotation", @@ -21,6 +22,7 @@ from diracx.db.sql import AuthDB as _AuthDB from diracx.db.sql import JobDB as _JobDB from diracx.db.sql import JobLoggingDB as _JobLoggingDB +from diracx.db.sql import ProxyDB as _ProxyDB from diracx.db.sql import SandboxMetadataDB as _SandboxMetadataDB from diracx.db.sql import TaskQueueDB as _TaskQueueDB @@ -36,6 +38,7 @@ def add_settings_annotation(cls: T) -> T: AuthDB = Annotated[_AuthDB, Depends(_AuthDB.transaction)] JobDB = Annotated[_JobDB, Depends(_JobDB.transaction)] JobLoggingDB = Annotated[_JobLoggingDB, Depends(_JobLoggingDB.transaction)] +ProxyDB = Annotated[_ProxyDB, Depends(_ProxyDB.transaction)] SandboxMetadataDB = Annotated[ _SandboxMetadataDB, Depends(_SandboxMetadataDB.transaction) ]