From b1331f9161a8255471b7986429a00028f27741c1 Mon Sep 17 00:00:00 2001
From: Jacob Tomlinson <jacobtomlinson@users.noreply.github.com>
Date: Thu, 27 Apr 2023 17:04:22 +0100
Subject: [PATCH] Reauthenticate (#37)

---
 kr8s/_api.py            | 35 ++++++++++++++++++-----------------
 kr8s/_auth.py           |  8 ++++++--
 kr8s/portforward.py     |  2 +-
 kr8s/tests/test_auth.py | 16 ++++++++++++++++
 4 files changed, 41 insertions(+), 20 deletions(-)

diff --git a/kr8s/_api.py b/kr8s/_api.py
index e37dc86..64c2686 100644
--- a/kr8s/_api.py
+++ b/kr8s/_api.py
@@ -111,23 +111,24 @@ async def call_api(
         parts.append(url)
         url = "/".join(parts)
 
-        if websocket:
-            async with self._session.ws_connect(
-                url=url,
-                ssl=self._sslcontext,
-                **kwargs,
-            ) as response:
-                yield response
-        else:
-            async with self._session.request(
-                method=method,
-                url=url,
-                ssl=self._sslcontext,
-                raise_for_status=raise_for_status,
-                **kwargs,
-            ) as response:
-                # TODO catch self.auth error and reauth a couple of times before giving up
-                yield response
+        call_method = self._session.ws_connect if websocket else self._session.request
+        kwargs.update(url=url, ssl=self._sslcontext)
+        if not websocket:
+            kwargs.update(method=method, raise_for_status=raise_for_status)
+
+        auth_attempts = 0
+        while True:
+            try:
+                async with call_method(**kwargs) as response:
+                    yield response
+            except aiohttp.ClientResponseError as e:
+                if e.status in (401, 403) and auth_attempts < 3:
+                    auth_attempts += 1
+                    self.auth.reauthenticate()
+                    continue
+                else:
+                    raise
+            break
 
     @contextlib.asynccontextmanager
     async def _get_kind(
diff --git a/kr8s/_auth.py b/kr8s/_auth.py
index 4e33420..7f30ce8 100644
--- a/kr8s/_auth.py
+++ b/kr8s/_auth.py
@@ -39,9 +39,13 @@ def __init__(
         ).expanduser()
         if url:
             self.server = url
-        if serviceaccount and not self.server:
+        self.reauthenticate()
+
+    def reauthenticate(self):
+        """Reauthenticate with the server."""
+        if self._serviceaccount and not self.server:
             self.load_service_account()
-        if kubeconfig is not False and not self.server:
+        if self._kubeconfig is not False and not self.server:
             self.load_kubeconfig()
         if not self.server:
             raise ValueError("Unable to find valid credentials")
diff --git a/kr8s/portforward.py b/kr8s/portforward.py
index ff24233..d1b867f 100644
--- a/kr8s/portforward.py
+++ b/kr8s/portforward.py
@@ -191,7 +191,7 @@ async def _ws_to_tcp(self, writer):
                         # Keep track of our channels. Could be useful later for listening to multiple ports.
                         channels.append(message.data[0])
                     else:
-                        if message.data[0] % 2 == 1:
+                        if message.data[0] % 2 == 1:  # pragma: no cover
                             # Odd channels are for errors.
                             raise ConnectionClosedError(message.data[1:].decode())
                         writer.write(message.data[1:])
diff --git a/kr8s/tests/test_auth.py b/kr8s/tests/test_auth.py
index 32a7d49..c8d3658 100644
--- a/kr8s/tests/test_auth.py
+++ b/kr8s/tests/test_auth.py
@@ -4,6 +4,7 @@
 import tempfile
 from pathlib import Path
 
+import aiohttp
 import pytest
 import yaml
 
@@ -43,6 +44,21 @@ async def test_kubeconfig(k8s_cluster):
     assert "major" in version
 
 
+async def test_reauthenticate(k8s_cluster):
+    kubernetes = kr8s.api(kubeconfig=k8s_cluster.kubeconfig_path)
+    kubernetes.auth.reauthenticate()
+    version = await kubernetes.version()
+    assert "major" in version
+
+
+async def test_bad_auth(serviceaccount):
+    (Path(serviceaccount) / "token").write_text("abc123")
+    kubernetes = kr8s.api(serviceaccount=serviceaccount, kubeconfig="/no/file/here")
+    serviceaccount = Path(serviceaccount)
+    with pytest.raises(aiohttp.ClientResponseError):
+        await kubernetes.version()
+
+
 async def test_url(kubectl_proxy):
     kubernetes = kr8s.api(url=kubectl_proxy)
     version = await kubernetes.version()