Fixes #10498 - Add support for nifi client certificate auth. (#10499)

* feat: Added nifi support for client cert auth

* feat: fix code linting

* feat: addressed comments for nifi connections

* feat: fixed linting

ingestion/src/metadata/ingestion/source/pipeline/nifi/client.py

@@ -32,19 +32,40 @@ class NifiClient:
     Wrapper on top of Nifi REST API
     """
 
+    # pylint: disable=too-many-arguments
     def __init__(
-        self, host_port: str, username: str, password: str, verify: bool = False
+        self,
+        host_port: str,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        ca_file_path: Optional[str] = None,
+        client_cert_path: Optional[str] = None,
+        client_key_path: Optional[str] = None,
+        verify: bool = False,
     ):
         self._token = None
         self._resources = None
 
+        self.content_headers = {"Content-Type": "application/x-www-form-urlencoded"}
+        self.api_endpoint = host_port + "/nifi-api"
         self.username = username
         self.password = password
-        self.verify = verify
-        self.api_endpoint = host_port + "/nifi-api"
 
-        self.content_headers = {"Content-Type": "application/x-www-form-urlencoded"}
-        self.headers = {"Authorization": f"Bearer {self.token}", **self.content_headers}
+        if all(setting for setting in [self.username, self.password]):
+            self.verify = verify
+            self.headers = {
+                "Authorization": f"Bearer {self.token}",
+                **self.content_headers,
+            }
+            self.data = f"username={self.username}&password={self.password}"
+            self.client_cert = None
+        else:
+            self.data = None
+            self.verify = ca_file_path if ca_file_path else False
+            self.client_cert = (client_cert_path, client_key_path)
+            self.headers = self.content_headers
+            access = self.get("access")
+            logger.debug(access)
 
     @property
     def token(self) -> str:
@@ -58,7 +79,7 @@ class NifiClient:
                     f"{self.api_endpoint}/access/token",
                     verify=self.verify,
                     headers=self.content_headers,
-                    data=f"username={self.username}&password={self.password}",
+                    data=self.data,
                     timeout=REQUESTS_TIMEOUT,
                 )
                 self._token = res.text
@@ -84,7 +105,10 @@ class NifiClient:
             self._resources = self.get(RESOURCES)  # API endpoint
 
         # Get the first `resources` key from the dict
-        return self._resources.get(RESOURCES)  # Dict key
+        try:
+            return self._resources.get(RESOURCES)  # Dict key
+        except AttributeError:
+            return []
 
     def get(self, path: str) -> Optional[Any]:
         """
@@ -96,6 +120,7 @@ class NifiClient:
                 verify=self.verify,
                 headers=self.headers,
                 timeout=REQUESTS_TIMEOUT,
+                cert=self.client_cert,
             )
 
             return res.json()

ingestion/src/metadata/ingestion/source/pipeline/nifi/connection.py

@@ -13,6 +13,7 @@
 Source connection handler
 """
 from metadata.generated.schema.entity.services.connections.pipeline.nifiConnection import (
+    BasicAuthentication,
     NifiConnection,
 )
 from metadata.ingestion.connections.test_connections import SourceConnectionException
@@ -23,11 +24,21 @@ def get_connection(connection: NifiConnection) -> NifiClient:
     """
     Create connection
     """
+    if isinstance(connection.nifiConfig, BasicAuthentication):
+        return NifiClient(
+            host_port=connection.hostPort,
+            username=connection.nifiConfig.username,
+            password=connection.nifiConfig.password.get_secret_value()
+            if connection.nifiConfig.password
+            else None,
+            verify=connection.nifiConfig.verifySSL,
+        )
+
     return NifiClient(
         host_port=connection.hostPort,
-        username=connection.username,
-        password=connection.password.get_secret_value(),
-        verify=connection.verifySSL,
+        ca_file_path=connection.nifiConfig.certificateAuthorityPath,
+        client_cert_path=connection.nifiConfig.clientCertificatePath,
+        client_key_path=connection.nifiConfig.clientkeyPath,
     )
 
 

ingestion/tests/unit/topology/pipeline/test_nifi.py

@@ -55,9 +55,11 @@ mock_nifi_config = {
             "config": {
                 "type": "Nifi",
                 "hostPort": "https://localhost:8443",
-                "username": "username",
-                "password": "password",
-                "verifySSL": False,
+                "nifiConfig": {
+                    "username": "username",
+                    "password": "password",
+                    "verifySSL": False,
+                },
             }
         },
         "sourceConfig": {"config": {"type": "PipelineMetadata"}},

openmetadata-spec/src/main/resources/json/schema/entity/services/connections/pipeline/nifiConnection.json

@@ -11,6 +11,52 @@
       "type": "string",
       "enum": ["Nifi"],
       "default": "Nifi"
+    },
+    "basicAuthentication": {
+      "title": "Username/Password Authentication",
+      "description": "username/password auth",
+      "properties": {
+        "username": {
+          "title": "Username",
+          "description": "Nifi user to authenticate to the API.",
+          "type": "string"
+        },
+        "password": {
+          "title": "Password",
+          "description": "Nifi password to authenticate to the API.",
+          "type": "string",
+          "format": "password"
+        },
+        "verifySSL": {
+          "title": "Verify SSL",
+          "description": "Boolean marking if we need to verify the SSL certs for Nifi. False by default.",
+          "type": "boolean",
+          "default": false
+        }
+      },
+      "additionalProperties": false
+    },
+    "clientCertificateAuthentication": {
+      "title": "Client Certificate Authentication",
+      "description": "client certificate auth",
+      "properties": {
+        "certificateAuthorityPath":{
+          "title":"Certificat Authority Path",
+          "description": "Path to the root CA certificate",
+          "type": "string"
+        },
+        "clientCertificatePath":{
+          "title":"Client Certificat",
+          "description": "Path to the client certificate",
+          "type": "string"
+        },
+        "clientkeyPath":{
+          "title":"Client Key",
+          "description": "Path to the client key",
+          "type": "string"
+        }
+      },
+      "additionalProperties": false
     }
   },
   "properties": {
@@ -27,22 +73,17 @@
       "type": "string",
       "format": "uri"
     },
-    "username": {
-      "title": "Username",
-      "description": "Nifi user to authenticate to the API.",
-      "type": "string"
-    },
-    "password": {
-      "title": "Password",
-      "description": "Nifi password to authenticate to the API.",
-      "type": "string",
-      "format": "password"
-    },
-    "verifySSL": {
-      "title": "Verify SSL",
-      "description": "Boolean marking if we need to verify the SSL certs for Nifi. False by default.",
-      "type": "boolean",
-      "default": false
+    "nifiConfig": {
+      "title": "Nifi Credentials Configuration",
+      "description": "We support username/password or client certificate authentication",
+      "oneOf": [
+        {
+          "$ref": "#/definitions/basicAuthentication"
+        },
+        {
+          "$ref": "#/definitions/clientCertificateAuthentication"
+        }
+      ]
     },
     "supportsMetadataExtraction": {
       "title": "Supports Metadata Extraction",
@@ -50,5 +91,5 @@
     }
   },
   "additionalProperties": false,
-  "required": ["hostPort", "username", "password"]
+  "required": ["hostPort", "nifiConfig"]
 }