Add network policies

src/charm.py

@@ -54,6 +54,7 @@
 from ops.model import ActiveStatus, BlockedStatus, MaintenanceStatus, ModelError, WaitingStatus
 from ops.pebble import Error, ExecError, Layer
 
+from k8s_network_policies import K8sNetworkPoliciesHandler, NetworkPoliciesHandlerError
 from kratos import KratosAPI
 
 if TYPE_CHECKING:
@@ -117,6 +118,7 @@ def __init__(self, *args: Any) -> None:
             port=KRATOS_PUBLIC_PORT,
             strip_prefix=True,
         )
+        self.network_policy_handler = K8sNetworkPoliciesHandler(self)
 
         self.database = DatabaseRequires(
             self,
@@ -151,9 +153,13 @@ def __init__(self, *args: Any) -> None:
         self.framework.observe(self.database.on.database_created, self._on_database_created)
         self.framework.observe(self.database.on.endpoints_changed, self._on_database_changed)
         self.framework.observe(self.admin_ingress.on.ready, self._on_admin_ingress_ready)
+        self.framework.observe(self.admin_ingress.on.ready, self._apply_network_policies)
         self.framework.observe(self.admin_ingress.on.revoked, self._on_ingress_revoked)
+        self.framework.observe(self.admin_ingress.on.revoked, self._apply_network_policies)
         self.framework.observe(self.public_ingress.on.ready, self._on_public_ingress_ready)
+        self.framework.observe(self.public_ingress.on.ready, self._apply_network_policies)
         self.framework.observe(self.public_ingress.on.revoked, self._on_ingress_revoked)
+        self.framework.observe(self.public_ingress.on.revoked, self._apply_network_policies)
         self.framework.observe(
             self.external_provider.on.client_config_changed, self._on_client_config_changed
         )
@@ -522,6 +528,21 @@ def _on_admin_ingress_ready(self, event: IngressPerAppReadyEvent) -> None:
         self._handle_status_update_config(event)
         self._update_kratos_endpoints_relation_data(event)
 
+    def _apply_network_policies(self, event: HookEvent) -> None:
+        if not self.unit.is_leader():
+            return
+
+        try:
+            self.network_policy_handler.apply_ingress_policy(
+                [
+                    ("admin", [self.admin_ingress.relation]),
+                    ("public", [self.public_ingress.relation]),
+                    (),
+                ]
+            )
+        except NetworkPoliciesHandlerError:
+            event.defer()
+
     def _on_public_ingress_ready(self, event: IngressPerAppReadyEvent) -> None:
         if self.unit.is_leader():
             logger.info("This app's public ingress URL: %s", event.url)

src/k8s_network_policies.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+# Copyright 2022 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+"""A helper class for managing kubernetes network policies."""
+
+import logging
+from typing import List, Optional, Tuple, Union
+
+from lightkube import ApiError, Client
+from lightkube.models.meta_v1 import LabelSelector, ObjectMeta
+from lightkube.models.networking_v1 import (
+    NetworkPolicyIngressRule,
+    NetworkPolicyPeer,
+    NetworkPolicyPort,
+    NetworkPolicySpec,
+)
+from lightkube.resources.networking_v1 import NetworkPolicy
+from ops.charm import CharmBase
+from ops.model import Relation
+
+logger = logging.getLogger(__name__)
+
+
+class NetworkPoliciesHandlerError(Exception):
+    """Applying the network policies failed."""
+
+
+IngressPolicyDefinition = Tuple[Union[str, int], List[Relation]]
+
+
+class K8sNetworkPoliciesHandler:
+    """A helper class for managing kubernetes network policies."""
+
+    def __init__(self, charm: CharmBase) -> None:
+        self._charm = charm
+        self.client = Client(field_manager=charm.app.name, namespace=charm.model.name)
+
+    @property
+    def policy_name(self) -> str:
+        """The default policy name that will be created."""
+        return f"{self._charm.app.name}-network-policy"
+
+    def apply_ingress_policy(
+        self, policies: IngressPolicyDefinition, name: Optional[str] = None
+    ) -> None:
+        """Apply an ingress network policy about a related application.
+
+        Policies can be defined for multiple ports at once to allow ingress traffic
+        from related applications
+
+        If no policies are defined then all ingress traffic will be denied.
+
+        Example usage:
+
+            policies = [("admin", [admin_ingress_relation]), (8080, [public_ingress_relation])]
+            network_policy_handler.apply_ingress_policy(policies)
+        """
+        if not name:
+            name = self.policy_name
+
+        ingress = []
+        for port, relations in policies:
+            selectors = [
+                NetworkPolicyPeer(
+                    podSelector=LabelSelector(
+                        matchLabels={"app.kubernetes.io/name": relation.app.name}
+                    )
+                )
+                for relation in relations
+                if relation and relation.app
+            ]
+            ingress.append(
+                NetworkPolicyIngressRule(
+                    from_=selectors,
+                    ports=[NetworkPolicyPort(port=port)],
+                ),
+            )
+
+        policy = NetworkPolicy(
+            metadata=ObjectMeta(name=name),
+            spec=NetworkPolicySpec(
+                podSelector=LabelSelector(
+                    matchLabels={"app.kubernetes.io/name": self._charm.app.name}
+                ),
+                policyTypes=["Ingress"],
+                ingress=ingress,
+            ),
+        )
+
+        try:
+            self.client.apply(
+                policy,
+                namespace=self._charm.model.name,
+            )
+        except ApiError as e:
+            if e.status.code == 403:
+                msg = f"Kubernetes resources patch failed: `juju trust` this application. {e}"
+            else:
+                msg = f"Kubernetes resources patch failed: {e}"
+            logger.error(msg)
+            raise NetworkPoliciesHandlerError()
+
+    def delete_network_policy(self, name: Optional[str] = None) -> None:
+        """Delete a network policy rule."""
+        if not name:
+            name = self.policy_name
+
+        try:
+            self.client.delete(name, namespace=self._charm.model.name)
+        except ApiError as e:
+            if e.status.code == 403:
+                msg = f"Kubernetes resources patch failed: `juju trust` this application. {e}"
+            else:
+                msg = f"Kubernetes resources patch failed: {e}"
+            logger.error(msg)
+            raise NetworkPoliciesHandlerError()

tests/integration/conftest.py

@@ -0,0 +1,11 @@
+import os
+
+from lightkube import Client, KubeConfig
+import pytest
+
+KUBECONFIG = os.environ.get("TESTING_KUBECONFIG", "~/.kube/config")
+
+
+@pytest.fixture(scope="module")
+def client() -> Client:
+    return Client(config=KubeConfig.from_file(KUBECONFIG))

tests/integration/test_charm.py

@@ -10,6 +10,8 @@
 import requests
 import yaml
 from pytest_operator.plugin import OpsTest
+from lightkube import Client
+from lightkube.resources.networking_v1 import NetworkPolicy
 
 logger = logging.getLogger(__name__)
 
@@ -71,7 +73,7 @@ async def test_build_and_deploy(ops_test: OpsTest) -> None:
         assert ops_test.model.applications[APP_NAME].units[0].workload_status == "active"
 
 
-async def test_ingress_relation(ops_test: OpsTest) -> None:
+async def test_ingress_relation(ops_test: OpsTest, client: Client) -> None:
     await ops_test.model.deploy(
         TRAEFIK,
         application_name=TRAEFIK_PUBLIC_APP,
@@ -94,6 +96,10 @@ async def test_ingress_relation(ops_test: OpsTest) -> None:
         timeout=1000,
     )
 
+    # Validate network policies are created when ingress is provided
+    policy = client.get(NetworkPolicy, "kratos-network-policy")
+    assert policy
+
 
 async def test_has_public_ingress(ops_test: OpsTest) -> None:
     # Get the traefik address and try to reach kratos

tests/unit/conftest.py

@@ -15,6 +15,11 @@
 from kratos import KratosAPI
 
 
+@pytest.fixture(autouse=True)
+def lk_client(mocker):
+    return mocker.patch("charm.Client", autospec=True)
+
+
 @pytest.fixture()
 def harness(mocked_kubernetes_service_patcher: MagicMock) -> Harness:
     harness = Harness(KratosCharm)