Move all Relationship references to single dataclass Relationship (#11)

* Move all Relationship references to single dataclass

* Fixes bug in `attack_paths.py` where POD was used
    instead of PODS (due to internals pluralizing resources)

* Address review: fix bug in var order for GRANT relationship creation

* Format code using `poetry run black .`

* CI fixes

---------

Co-authored-by: Mohit Gupta <[email protected]>

Author: gsalaz98

icekube/attack_paths.py

@@ -2,6 +2,8 @@
 
 from typing import List
 
+from icekube.relationships import Relationship
+
 WORKLOAD_TYPES = [
     "ReplicationController",
     "DaemonSet",
@@ -27,40 +29,40 @@ def workload_query(
 
 attack_paths = {
     # Subject -> Role Bindings
-    "BOUND_TO": "MATCH (src)-[:BOUND_TO]->(dest)",
+    Relationship.BOUND_TO: "MATCH (src)-[:BOUND_TO]->(dest)",
     # Role Binding -> Role
-    "GRANTS_PERMISSION": "MATCH (src)-[:GRANTS_PERMISSION]->(dest)",
+    Relationship.GRANTS_PERMISSION: "MATCH (src)-[:GRANTS_PERMISSION]->(dest)",
     # Pod -> Service Account
-    "USES_ACCOUNT": "MATCH (src:Pod)-[:USES_ACCOUNT]->(dest:ServiceAccount)",
+    Relationship.USES_ACCOUNT: "MATCH (src:Pod)-[:USES_ACCOUNT]->(dest:ServiceAccount)",
     # Pod -> Secrett
-    "MOUNTS_SECRET": "MATCH (src:Pod)-[:MOUNTS_SECRET]->(dest:Secret)",
+    Relationship.MOUNTS_SECRET: "MATCH (src:Pod)-[:MOUNTS_SECRET]->(dest:Secret)",
     # Subject has permission to create pod within namespace with target
     # Service Account
-    "CREATE_POD_WITH_SA": f"""
+    Relationship.CREATE_POD_WITH_SA: f"""
         MATCH (src)-[:GRANTS_PODS_CREATE|{create_workload_query()}]->(ns:Namespace)<-[:WITHIN_NAMESPACE]-(dest:ServiceAccount)
         """,
     # Subject has permission to update workload within namespace with target
     # Service Account
-    "UPDATE_WORKLOAD_WITH_SA": f"""
+    Relationship.UPDATE_WORKLOAD_WITH_SA: f"""
         MATCH (src)-[:GRANTS_UPDATE|GRANTS_PATCH]->(workload)-[:WITHIN_NAMESPACE]->(ns:Namespace)<-[:WITHIN_NAMESPACE]-(dest:ServiceAccount)
         WHERE {workload_query()}
         """,
     # Subject -> Pod
-    "EXEC_INTO": "MATCH (src)-[:GRANTS_EXEC_CREATE]->(dest:Pod)<-[:GRANTS_GET]-(src)",
+    Relationship.EXEC_INTO: "MATCH (src)-[:GRANTS_EXEC_CREATE]->(dest:Pod)<-[:GRANTS_GET]-(src)",
     # Subject -> Pod
-    "REPLACE_IMAGE": "MATCH (src)-[:GRANTS_PATCH]->(dest:Pod)",
+    Relationship.REPLACE_IMAGE: "MATCH (src)-[:GRANTS_PATCH]->(dest:Pod)",
     # Subject -> Pod
-    "DEBUG_POD": "MATCH (src)-[:GRANTS_EPHEMERAL_PATCH]->(dest:Pod)",
+    Relationship.DEBUG_POD: "MATCH (src)-[:GRANTS_EPHEMERAL_PATCH]->(dest:Pod)",
     # Subject has permission to read authentication token for Service Account
-    "GET_AUTHENTICATION_TOKEN_FOR": """
+    Relationship.GET_AUTHENTICATION_TOKEN_FOR: """
         MATCH (src)-[:GRANTS_GET|GRANTS_LIST|GRANTS_WATCH]->(secret:Secret)-[:AUTHENTICATION_TOKEN_FOR]->(dest:ServiceAccount)
         """,
     # Subject -> Secret
-    "ACCESS_SECRET": "MATCH (src)-[:GRANTS_GET|GRANTS_LIST|GRANTS_WATCH]->(dest:Secret)",
+    Relationship.ACCESS_SECRET: "MATCH (src)-[:GRANTS_GET|GRANTS_LIST|GRANTS_WATCH]->(dest:Secret)",
     # Generate service account token
-    "GENERATE_TOKEN": "MATCH (src)-[:GRANTS_TOKEN_CREATE]->(dest:ServiceAccount)",
+    Relationship.GENERATE_TOKEN: "MATCH (src)-[:GRANTS_TOKEN_CREATE]->(dest:ServiceAccount)",
     # RBAC escalate verb to change a role to be more permissive
-    "RBAC_ESCALATE_TO": [
+    Relationship.RBAC_ESCALATE_TO: [
         # RoleBindings
         """
         MATCH (src:RoleBinding)-[:GRANTS_ESCALATE]->(role)-[:WITHIN_NAMESPACE]->(:Namespace)<-[:WITHIN_NAMESPACE]-(dest)
@@ -73,39 +75,39 @@ def workload_query(
         """,
     ],
     # Subject -> User / Group / ServiceAccount
-    "GENERATE_CLIENT_CERTIFICATE": """
+    Relationship.GENERATE_CLIENT_CERTIFICATE: """
         MATCH (src)-[:GRANTS_CERTIFICATESIGNINGREQUESTS_CREATE]->(cluster:Cluster), (dest)
         WHERE (src)-[:HAS_CSR_APPROVAL]->(cluster) AND (src)-[:GRANTS_APPROVE]->(:Signer {
           name: "kubernetes.io/kube-apiserver-client"
         }) AND (dest:User OR dest:Group OR dest:ServiceAccount)
         """,
     # Impersonate
-    "CAN_IMPERSONATE": "MATCH (src)-[:GRANTS_IMPERSONATE]->(dest)",
+    Relationship.CAN_IMPERSONATE: "MATCH (src)-[:GRANTS_IMPERSONATE]->(dest)",
     # Pod breakout
-    "IS_PRIVILEGED": "MATCH (src:Pod {privileged: true})<-[:HOSTS_POD]-(dest:Node)",
-    "CAN_CGROUP_BREAKOUT": 'MATCH (src:Pod)<-[:HOSTS_POD]-(dest:Node) WHERE "SYS_ADMIN" in src.capabilities',
-    "CAN_LOAD_KERNEL_MODULES": 'MATCH (src:Pod)<-[:HOSTS_POD]-(dest:Node) WHERE "SYS_MODULE" in src.capabilities',
-    "CAN_ACCESS_DANGEROUS_HOST_PATH": "MATCH (src:Pod {dangerous_host_path: true})<-[:HOSTS_POD]-(dest:Node)",
-    "CAN_NSENTER_HOST": 'MATCH (src:Pod {hostPID: true})<-[:HOSTS_POD]-(dest:Node) WHERE all(x in ["SYS_ADMIN", "SYS_PTRACE"] WHERE x in src.capabilities)',
-    "CAN_ACCESS_HOST_FD": 'MATCH (src:Pod)<-[:HOSTS_POD]-(dest:Node) WHERE "DAC_READ_SEARCH" in src.capabilities',
+    Relationship.IS_PRIVILEGED: "MATCH (src:Pod {privileged: true})<-[:HOSTS_POD]-(dest:Node)",
+    Relationship.CAN_CGROUP_BREAKOUT: 'MATCH (src:Pod)<-[:HOSTS_POD]-(dest:Node) WHERE "SYS_ADMIN" in src.capabilities',
+    Relationship.CAN_LOAD_KERNEL_MODULES: 'MATCH (src:Pod)<-[:HOSTS_POD]-(dest:Node) WHERE "SYS_MODULE" in src.capabilities',
+    Relationship.CAN_ACCESS_DANGEROUS_HOST_PATH: "MATCH (src:Pod {dangerous_host_path: true})<-[:HOSTS_POD]-(dest:Node)",
+    Relationship.CAN_NSENTER_HOST: 'MATCH (src:Pod {hostPID: true})<-[:HOSTS_POD]-(dest:Node) WHERE all(x in ["SYS_ADMIN", "SYS_PTRACE"] WHERE x in src.capabilities)',
+    Relationship.CAN_ACCESS_HOST_FD: 'MATCH (src:Pod)<-[:HOSTS_POD]-(dest:Node) WHERE "DAC_READ_SEARCH" in src.capabilities',
     # Can jump to pods running on node
-    "ACCESS_POD": "MATCH (src:Node)-[:HOSTS_POD]->(dest:Pod)",
+    Relationship.ACCESS_POD: "MATCH (src:Node)-[:HOSTS_POD]->(dest:Pod)",
     # Can exec into pods on a node
-    "CAN_EXEC_THROUGH_KUBELET": "MATCH (src)-[:GRANTS_PROXY_CREATE]->(:Node)-[:HOSTS_POD]->(dest:Pod)",
+    Relationship.CAN_EXEC_THROUGH_KUBELET: "MATCH (src)-[:GRANTS_PROXY_CREATE]->(:Node)-[:HOSTS_POD]->(dest:Pod)",
     # Can update aws-auth ConfigMap
-    "UPDATE_AWS_AUTH": """
+    Relationship.UPDATE_AWS_AUTH: """
         MATCH (src)-[:GRANTS_PATCH|GRANTS_UPDATE]->(:ConfigMap {
           name: 'aws-auth', namespace: 'kube-system'
         }), (dest:Group {
           name: 'system:masters'
         })
         """,
-    "AZURE_POD_IDENTITY_EXCEPTION": [
+    Relationship.AZURE_POD_IDENTITY_EXCEPTION: [
         # Create workload based of existing APIE
         f"""
         MATCH (src)-[:GRANTS_GET|GRANTS_LIST|GRANTS_WATCH]->(azexc:AzurePodIdentityException)-[:WITHIN_NAMESPACE]->(ns:Namespace), (dest:ClusterRoleBinding)
         WHERE (dest.name = 'aks-cluster-admin-binding' OR dest.name = 'aks-cluster-admin-binding-aad') AND (EXISTS {{
-            MATCH (src)-[:{create_workload_query()}|GRANTS_POD_CREATE]->(ns)
+            MATCH (src)-[:{create_workload_query()}|GRANTS_PODS_CREATE]->(ns)
         }} OR EXISTS {{
             MATCH (src)-[:GRANTS_PATCH|GRANTS_UPDATE]->(workload)-[:WITHIN_NAMESPACE]->(ns)
             WHERE {workload_query()}

icekube/models/base.py

@@ -5,6 +5,7 @@
 import traceback
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
+from icekube.relationships import Relationship
 from icekube.utils import to_camel_case
 from kubernetes import client
 from pydantic import BaseModel, Field, root_validator
@@ -214,7 +215,7 @@ def relationships(
             relationships += [
                 (
                     self,
-                    "WITHIN_NAMESPACE",
+                    Relationship.WITHIN_NAMESPACE,
                     ns,
                 ),
             ]

icekube/models/clusterrolebinding.py

@@ -10,6 +10,7 @@
 from icekube.models.serviceaccount import ServiceAccount
 from icekube.models.user import User
 from icekube.neo4j import find_or_mock, get_cluster_object, mock
+from icekube.relationships import Relationship
 from pydantic import root_validator
 from pydantic.fields import Field
 
@@ -82,14 +83,16 @@ def relationships(
         initial: bool = True,
     ) -> List[RELATIONSHIP]:
         relationships = super().relationships()
-        relationships += [(self, "GRANTS_PERMISSION", self.role)]
-        relationships += [(subject, "BOUND_TO", self) for subject in self.subjects]
+        relationships += [(self, Relationship.GRANTS_PERMISSION, self.role)]
+        relationships += [
+            (subject, Relationship.BOUND_TO, self) for subject in self.subjects
+        ]
 
         if not initial:
             for role_rule in self.role.rules:
                 if role_rule.contains_csr_approval:
                     relationships.append(
-                        (self, "HAS_CSR_APPROVAL", get_cluster_object()),
+                        (self, Relationship.HAS_CSR_APPROVAL, get_cluster_object()),
                     )
                 for relationship, resource in role_rule.affected_resource_query():
                     relationships.append((self, relationship, resource))

icekube/models/pod.py

@@ -10,6 +10,7 @@
 from icekube.models.secret import Secret
 from icekube.models.serviceaccount import ServiceAccount
 from icekube.neo4j import mock
+from icekube.relationships import Relationship
 from pydantic import root_validator
 
 CAPABILITIES = [
@@ -249,14 +250,14 @@ def relationships(
         relationships = super().relationships()
 
         if self.service_account:
-            relationships += [(self, "USES_ACCOUNT", self.service_account)]
+            relationships += [(self, Relationship.USES_ACCOUNT, self.service_account)]
         if self.node:
-            relationships += [(self.node, "HOSTS_POD", self)]
+            relationships += [(self.node, Relationship.HOSTS_POD, self)]
         for secret in self.mounted_secrets:
             relationships += [
                 (
                     self,
-                    "MOUNTS_SECRET",
+                    Relationship.MOUNTS_SECRET,
                     mock(Secret, namespace=cast(str, self.namespace), name=secret),
                 ),
             ]

icekube/models/policyrule.py

@@ -3,6 +3,7 @@
 from fnmatch import fnmatch
 from typing import Dict, Iterator, List, Optional, Tuple, Union
 
+from icekube.relationships import Relationship
 from pydantic import BaseModel
 from pydantic.fields import Field
 
@@ -91,23 +92,22 @@ def affected_resource_query(
                 else:
                     query_filter = {"kind": "Cluster"}
                     yield (
-                        f"GRANTS_{resource}_CREATE".upper().replace("-", "_"),
+                        Relationship.generate_grant("CREATE", resource),
                         generate_query(query_filter),
                     )
                     query_filter = {"kind": "Namespace"}
                 yield (
-                    f"GRANTS_{resource}_CREATE".upper().replace("-", "_"),
+                    Relationship.generate_grant("CREATE", resource),
                     generate_query(query_filter),
                 )
                 valid_verbs.remove("create")
 
             if not valid_verbs:
                 continue
 
-            if sub_resource is None:
-                tags = [f"GRANTS_{verb}".upper() for verb in valid_verbs]
-            else:
-                tags = [f"GRANTS_{sub_resource}_{verb}".upper() for verb in valid_verbs]
+            tags = [
+                Relationship.generate_grant(verb, sub_resource) for verb in valid_verbs
+            ]
 
             if not self.resourceNames:
                 yield (tags, generate_query(find_filter))

icekube/models/rolebinding.py

@@ -10,6 +10,7 @@
 from icekube.models.role import Role
 from icekube.models.serviceaccount import ServiceAccount
 from icekube.models.user import User
+from icekube.relationships import Relationship
 from pydantic import root_validator
 from pydantic.fields import Field
 
@@ -39,8 +40,10 @@ def relationships(
         initial: bool = True,
     ) -> List[RELATIONSHIP]:
         relationships = super().relationships()
-        relationships += [(self, "GRANTS_PERMISSION", self.role)]
-        relationships += [(subject, "BOUND_TO", self) for subject in self.subjects]
+        relationships += [(self, Relationship.GRANTS_PERMISSION, self.role)]
+        relationships += [
+            (subject, Relationship.BOUND_TO, self) for subject in self.subjects
+        ]
 
         if not initial:
             for role_rule in self.role.rules:

icekube/models/secret.py

@@ -5,6 +5,7 @@
 
 from icekube.models.base import RELATIONSHIP, Resource
 from icekube.neo4j import mock
+from icekube.relationships import Relationship
 from pydantic import root_validator
 
 
@@ -52,7 +53,7 @@ def relationships(self, initial: bool = True) -> List[RELATIONSHIP]:
                 relationships.append(
                     (
                         self,
-                        "AUTHENTICATION_TOKEN_FOR",
+                        Relationship.AUTHENTICATION_TOKEN_FOR,
                         account,
                     ),
                 )

icekube/models/serviceaccount.py

@@ -6,6 +6,7 @@
 from icekube.models.base import RELATIONSHIP, Resource
 from icekube.models.secret import Secret
 from icekube.neo4j import mock
+from icekube.relationships import Relationship
 from pydantic import root_validator
 from pydantic.fields import Field
 
@@ -39,5 +40,7 @@ def relationships(
         initial: bool = True,
     ) -> List[RELATIONSHIP]:
         relationships = super().relationships()
-        relationships += [(x, "AUTHENTICATION_TOKEN_FOR", self) for x in self.secrets]
+        relationships += [
+            (x, Relationship.AUTHENTICATION_TOKEN_FOR, self) for x in self.secrets
+        ]
         return relationships

icekube/relationships.py

@@ -0,0 +1,101 @@
+from typing import ClassVar, Optional
+
+
+class Relationship:
+    """Consolidates the various relationship types into a single class.
+
+    This allows for better tracking of where we assign each relationship
+    across the codebase.
+
+    Relationships in the order (ObjectOne, RELATIONSHIP, ObjectTwo) are
+    in this direction in neo4j: (ObjectOne)-[:RELATIONSHIP]->(ObjectTwo)
+    """
+
+    HOSTS_POD: ClassVar[str] = "HOSTS_POD"
+
+    AUTHENTICATION_TOKEN_FOR: ClassVar[str] = "AUTHENTICATION_TOKEN_FOR"
+    GET_AUTHENTICATION_TOKEN_FOR: ClassVar[str] = "GET_AUTHENTICATION_TOKEN_FOR"
+
+    WITHIN_NAMESPACE: ClassVar[str] = "WITHIN_NAMESPACE"
+
+    GRANTS_PODS_CREATE: ClassVar[str] = "GRANTS_PODS_CREATE"
+    GRANTS_REPLICATIONCONTROLLERS_CREATE: ClassVar[
+        str
+    ] = "GRANTS_REPLICATIONCONTROLLERS_CREATE"
+    GRANTS_DAEMONSETS_CREATE: ClassVar[str] = "GRANTS_DAEMONSETS_CREATE"
+    GRANTS_DEPLOYMENTS_CREATE: ClassVar[str] = "GRANTS_DEPLOYMENTS_CREATE"
+    GRANTS_REPLICASETS_CREATE: ClassVar[str] = "GRANTS_REPLICASETS_CREATE"
+    GRANTS_STATEFULSETS_CREATE: ClassVar[str] = "GRANTS_STATEFULSETS_CREATE"
+    GRANTS_CRONJOBS_CREATE: ClassVar[str] = "GRANTS_CRONJOBS_CREATE"
+    GRANTS_JOBS_CREATE: ClassVar[str] = "GRANTS_JOBS_CREATE"
+
+    GRANTS_AZUREPODIDENTITYEXCEPTIONS_CREATE: ClassVar[
+        str
+    ] = "GRANTS_AZUREPODIDENTITYEXCEPTIONS_CREATE"
+    GRANTS_CERTIFICATESIGNINGREQUESTS_CREATE: ClassVar[
+        str
+    ] = "GRANTS_CERTIFICATESIGNINGREQUESTS_CREATE"
+    GRANTS_PROXY_CREATE: ClassVar[str] = "GRANTS_PROXY_CREATE"
+
+    GRANTS_GET: ClassVar[str] = "GRANTS_GET"
+    GRANTS_LIST: ClassVar[str] = "GRANTS_LIST"
+    GRANTS_UPDATE: ClassVar[str] = "GRANTS_UPDATE"
+    GRANTS_WATCH: ClassVar[str] = "GRANTS_WATCH"
+    GRANTS_PATCH: ClassVar[str] = "GRANTS_PATCH"
+    GRANTS_APPROVE: ClassVar[str] = "GRANTS_APPROVE"
+    GRANTS_PERMISSION: ClassVar[str] = "GRANTS_PERMISSION"
+
+    GRANTS_ESCALATE: ClassVar[str] = "GRANTS_ESCALATE"
+    GRANTS_IMPERSONATE: ClassVar[str] = "GRANTS_IMPERSONATE"
+    GRANTS_TOKEN_CREATE: ClassVar[str] = "GRANTS_TOKEN_CREATE"
+    GRANTS_EPHEMERAL_PATCH: ClassVar[str] = "GRANTS_EPHEMERAL_PATCH"
+
+    BOUND_TO: ClassVar[str] = "BOUND_TO"
+    USES_ACCOUNT: ClassVar[str] = "USES_ACCOUNT"
+    MOUNTS_SECRET: ClassVar[str] = "MOUNTS_SECRET"
+    CREATE_POD_WITH_SA: ClassVar[str] = "CREATE_POD_WITH_SA"
+    UPDATE_WORKLOAD_WITH_SA: ClassVar[str] = "UPDATE_WORKLOAD_WITH_SA"
+
+    EXEC_INTO: ClassVar[str] = "EXEC_INTO"
+    REPLACE_IMAGE: ClassVar[str] = "REPLACE_IMAGE"
+    DEBUG_POD: ClassVar[str] = "DEBUG_POD"
+
+    ACCESS_SECRET: ClassVar[str] = "ACCESS_SECRET"
+    GENERATE_TOKEN: ClassVar[str] = "GENERATE_TOKEN"
+    RBAC_ESCALATE_TO: ClassVar[str] = "RBAC_ESCALATE_TO"
+
+    GENERATE_CLIENT_CERTIFICATE: ClassVar[str] = "GENERATE_CLIENT_CERTIFICATE"
+    HAS_CSR_APPROVAL: ClassVar[str] = "HAS_CSR_APPROVAL"
+
+    CAN_IMPERSONATE: ClassVar[str] = "CAN_IMPERSONATE"
+
+    IS_PRIVILEGED: ClassVar[str] = "IS_PRIVILEGED"
+    CAN_CGROUP_BREAKOUT: ClassVar[str] = "CAN_CGROUP_BREAKOUT"
+    CAN_LOAD_KERNEL_MODULES: ClassVar[str] = "CAN_LOAD_KERNEL_MODULES"
+    CAN_ACCESS_DANGEROUS_HOST_PATH: ClassVar[str] = "CAN_ACCESS_DANGEROUS_HOST_PATH"
+    CAN_NSENTER_HOST: ClassVar[str] = "CAN_NSENTER_HOST"
+    CAN_ACCESS_HOST_FD: ClassVar[str] = "CAN_ACCESS_HOST_FD"
+    CAN_EXEC_THROUGH_KUBELET: ClassVar[str] = "CAN_EXEC_THROUGH_KUBELET"
+
+    ACCESS_POD: ClassVar[str] = "ACCESS_POD"
+    UPDATE_AWS_AUTH: ClassVar[str] = "UPDATE_AWS_AUTH"
+    AZURE_POD_IDENTITY_EXCEPTION: ClassVar[str] = "AZURE_POD_IDENTITY_EXCEPTION"
+
+    # Current resource defines the spec/creation of the subresource
+    DEFINES: ClassVar[str] = "DEFINES"
+    # Defines a reference to another object (e.g. Pod -> ServiceAccount)
+    REFERENCES: ClassVar[str] = "REFERENCES"
+    # Directly consumes a resource (e.g. PersistentVolumeClaim -> PersistentVolume)
+    CONSUMES: ClassVar[str] = "CONSUMES"
+    # Indirectly consumes a resource, without an exclusive relationship to the refering
+    # node (e.g. PersistentVolume -> StorageClass)
+    USES: ClassVar[str] = "USES"
+    # Defines ownership of a resource (e.g. Deployment-[:OWNS]->ReplicaSet)
+    OWNS: ClassVar[str] = "OWNS"
+
+    @staticmethod
+    def generate_grant(verb: str, sub_resource: Optional[str]) -> str:
+        if sub_resource is None:
+            return f"GRANTS_{verb.upper()}".replace("-", "_")
+
+        return f"GRANTS_{sub_resource}_{verb}".upper().replace("-", "_")