From ed094e11ec59409c6cb361fa871e9b5e3da02172 Mon Sep 17 00:00:00 2001
From: Christian Heimes <cheimes@redhat.com>
Date: Oct 02 2023 21:40:57 +0000
Subject: Add context manager to ipalib.API


`ipalib.API` instances like `ipalib.api` now provide a context manager
that connects and disconnects the API object. Users no longer have to
deal with different types of backends or finalize the API correctly.

```python
import ipalib

with ipalib.api as api:
    api.Commands.ping()
```

See: https://pagure.io/freeipa/issue/9443
Signed-off-by: Christian Heimes <cheimes@redhat.com>
Reviewed-By: Alexander Bokovoy <abokovoy@redhat.com>

---

diff --git a/ipalib/__init__.py b/ipalib/__init__.py
index ef4839c..906e026 100644
--- a/ipalib/__init__.py
+++ b/ipalib/__init__.py
@@ -936,6 +936,68 @@ Registry = plugable.Registry
 class API(plugable.API):
     bases = (Command, Object, Method, Backend, Updater)
 
+    def __enter__(self):
+        """Context manager for IPA API
+
+        The context manager connects the backend connect on enter and
+        disconnects on exit. The process must have access to a valid Kerberos
+        ticket or have automatic authentication with a keytab or gssproxy
+        set up. The connection type depends on ``in_server`` and ``context``
+        options. Server connections use LDAP while clients use JSON-RPC over
+        HTTPS.
+
+        The context manager also finalizes the API object, in case it hasn't
+        been finalized yet. It is possible to use a custom API object. In
+        that case, the global API object must be finalized, first. Some
+        options like logging only apply to global ``ipalib.api`` object.
+
+        Usage with global api object::
+
+            import os
+            import ipalib
+
+            # optional: automatic authentication with a KRB5 keytab
+            os.environ.update(
+                KRB5_CLIENT_KTNAME="/path/to/service.keytab",
+                KRB5RCACHENAME="FILE:/path/to/tmp/service.ccache",
+            )
+
+            # optional: override settings (once per process)
+            overrides = {}
+            ipalib.api.bootstrap(**overrides)
+
+            with ipalib.api as api:
+                host = api.Command.host_show(api.env.host)
+                user = api.Command.user_show("admin")
+
+        """
+        # Several IPA module require api.env at import time, some even
+        # a fully finalized ipalib.ap, e.g. register() with MethodOverride.
+        if self is not api and not api.isdone("finalize"):
+            raise RuntimeError("global ipalib.api must be finalized first.")
+        # initialize this api
+        if not self.isdone("finalize"):
+            self.finalize()
+        # connect backend, server and client use different backends.
+        if self.env.in_server:
+            conn = self.Backend.ldap2
+        else:
+            conn = self.Backend.rpcclient
+        if conn.isconnected():
+            raise RuntimeError("API is already connected")
+        else:
+            conn.connect()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Disconnect backend on exit"""
+        if self.env.in_server:
+            conn = self.Backend.ldap2
+        else:
+            conn = self.Backend.rpcclient
+        if conn.isconnected():
+            conn.disconnect()
+
     @property
     def packages(self):
         if self.env.in_server:
diff --git a/ipatests/test_integration/example_cli.py b/ipatests/test_integration/example_cli.py
new file mode 100644
index 0000000..bd85ac3
--- /dev/null
+++ b/ipatests/test_integration/example_cli.py
@@ -0,0 +1,19 @@
+import os
+
+import ipalib
+from ipaplatform.paths import paths
+
+# authenticate with host keytab and custom ccache
+os.environ.update(
+    KRB5_CLIENT_KTNAME=paths.KRB5_KEYTAB,
+)
+
+# custom options
+overrides = {"context": "example_cli"}
+ipalib.api.bootstrap(**overrides)
+
+with ipalib.api as api:
+    user = api.Command.user_show("admin")
+    print(user)
+
+assert not api.Backend.rpcclient.isconnected()
diff --git a/ipatests/test_integration/test_commands.py b/ipatests/test_integration/test_commands.py
index b6ad243..70e118d 100644
--- a/ipatests/test_integration/test_commands.py
+++ b/ipatests/test_integration/test_commands.py
@@ -13,6 +13,7 @@ import random
 import shlex
 import ssl
 from itertools import chain, repeat
+import sys
 import textwrap
 import time
 import pytest
@@ -1557,6 +1558,34 @@ class TestIPACommand(IntegrationTest):
 
         assert 'Discovered server %s' % self.master.hostname in result
 
+    def test_ipa_context_manager(self):
+        """Exercise ipalib.api context manager and KRB5_CLIENT_KTNAME auth
+
+        The example_cli.py script uses the context manager to connect and
+        disconnect the global ipalib.api object. The test also checks whether
+        KRB5_CLIENT_KTNAME env var automatically acquires a TGT.
+        """
+        host = self.clients[0]
+        tasks.kdestroy_all(host)
+
+        here = os.path.abspath(os.path.dirname(__file__))
+        with open(os.path.join(here, "example_cli.py")) as f:
+            contents = f.read()
+
+        # upload script and run with Python executable
+        script = "/tmp/example_cli.py"
+        host.put_file_contents(script, contents)
+        result = host.run_command([sys.executable, script])
+
+        # script prints admin account
+        admin_princ = f"admin@{host.domain.realm}"
+        assert admin_princ in result.stdout_text
+
+        # verify that auto-login did use correct principal
+        host_princ = f"host/{host.hostname}@{host.domain.realm}"
+        result = host.run_command([paths.KLIST])
+        assert host_princ in result.stdout_text
+
 
 class TestIPACommandWithoutReplica(IntegrationTest):
     """