From 2a4f7824d84ef9f8a7d2ca94934ef62e7f2c0868 Mon Sep 17 00:00:00 2001
From: Nikos <nikos.sklikas@canonical.com>
Date: Thu, 20 Apr 2023 11:46:20 +0300
Subject: [PATCH] Working version

This is fairly complex, will probably have to refactor it
---
 src/charm.py                    | 12 +++++++++---
 src/k8s_network_policies.py     | 23 +++++++++++++++++++----
 tests/integration/test_charm.py |  2 +-
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 50dbf79c..3b7cb09b 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -54,7 +54,11 @@
 from ops.model import ActiveStatus, BlockedStatus, MaintenanceStatus, ModelError, WaitingStatus
 from ops.pebble import Error, ExecError, Layer
 
-from k8s_network_policies import K8sNetworkPoliciesHandler, NetworkPoliciesHandlerError
+from k8s_network_policies import (
+    K8sNetworkPoliciesHandler,
+    NetworkPoliciesHandlerError,
+    PortDefinition,
+)
 from kratos import KratosAPI
 
 if TYPE_CHECKING:
@@ -535,8 +539,10 @@ def _apply_network_policies(self, event: HookEvent) -> None:
         try:
             self.network_policy_handler.apply_ingress_policy(
                 [
-                    ("admin", [self.admin_ingress.relation]),
-                    ("public", [self.public_ingress.relation]),
+                    (PortDefinition(1, KRATOS_PUBLIC_PORT - 1), ()),
+                    (PortDefinition(KRATOS_PUBLIC_PORT), [self.public_ingress.relation]),
+                    (PortDefinition(KRATOS_ADMIN_PORT), [self.admin_ingress.relation]),
+                    (PortDefinition(KRATOS_ADMIN_PORT + 1, 65535), ()),
                 ]
             )
         except NetworkPoliciesHandlerError:
diff --git a/src/k8s_network_policies.py b/src/k8s_network_policies.py
index 132e9076..db760f58 100644
--- a/src/k8s_network_policies.py
+++ b/src/k8s_network_policies.py
@@ -5,6 +5,7 @@
 """A helper class for managing kubernetes network policies."""
 
 import logging
+from dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 
 from lightkube import ApiError, Client
@@ -26,7 +27,22 @@ class NetworkPoliciesHandlerError(Exception):
     """Applying the network policies failed."""
 
 
-IngressPolicyDefinition = Tuple[Union[str, int], List[Relation]]
+@dataclass
+class PortDefinition:
+    """Network Policy port definition."""
+
+    port: Union[str, int]
+    end_port: Optional[int] = None
+    protocol: Optional[str] = "TCP"
+
+    def to_resource(self):
+        """Convert class to NetworkPolicyPort."""
+        if not self.end_port:
+            return NetworkPolicyPort(port=self.port, protocol=self.protocol)
+        return NetworkPolicyPort(port=self.port, endPort=self.end_port, protocol=self.protocol)
+
+
+IngressPolicyDefinition = Tuple[PortDefinition, List[Relation]]
 
 
 class K8sNetworkPoliciesHandler:
@@ -73,7 +89,7 @@ def apply_ingress_policy(
             ingress.append(
                 NetworkPolicyIngressRule(
                     from_=selectors,
-                    ports=[NetworkPolicyPort(port=port)],
+                    ports=[port.to_resource()],
                 ),
             )
 
@@ -86,9 +102,8 @@ def apply_ingress_policy(
                         "kubernetes.io/metadata.name": self._charm.model.name,
                     }
                 ),
-                policyTypes=["Ingress", "Egress"],
+                policyTypes=["Ingress"],
                 ingress=ingress,
-                egress=[{}]
             ),
         )
 
diff --git a/tests/integration/test_charm.py b/tests/integration/test_charm.py
index 72869465..e18b26ca 100644
--- a/tests/integration/test_charm.py
+++ b/tests/integration/test_charm.py
@@ -97,7 +97,7 @@ async def test_ingress_relation(ops_test: OpsTest, client: Client) -> None:
     )
 
     # Validate network policies are created when ingress is provided
-    policy = client.get(NetworkPolicy, "kratos-network-policy")
+    policy = client.get(NetworkPolicy, "kratos-network-policy", namespace=ops_test.model.name)
     assert policy