Management of NiFi accounts + lot of refactoring

config.py

+# Various constants and parameters
+
+NIFI_CONTAINERS = [
+    'soctools-nifi-1',
+    'soctools-nifi-2',
+    'soctools-nifi-3',
+]

main.py

@@ -3,7 +3,6 @@ import sys
 from datetime import datetime, timezone
 import os.path
 import re
-import subprocess
 from typing import List, Dict, Optional
 
 from flask import Flask, render_template, request, make_response, redirect, flash
@@ -15,6 +14,9 @@ import requests
 import yaml
 from dataclasses import dataclass, field # external package - backport of Py3.7 dataclasses to Py3.6
 
+from nifi import *
+
+
 app = Flask(__name__)
 app.secret_key = "ASDF1234 - CHANGE ME!"
 
@@ -60,6 +62,7 @@ class UserAccount:
     dn: str
     kcid: Optional[str] = field(default=None) # keycloak ID
     ts_created: Optional[datetime] = field(default=None) # timezone-aware datetime in UTC
+    components: Optional[Dict[str, bool]] = field(default_factory=dict) # Presence of the account in SOCtools components that don't use Keycloak directly (NiFi, MISP, TheHive, ...)
 
     def to_keycloak_representation(self) -> Dict:
         """
@@ -134,7 +137,7 @@ def kc_get_users() -> List[UserAccount]:
     except (ValueError, AssertionError):
         raise KeycloakError("Can't get list of users: Unexpected content of response from Keycloak")
 
-def kc_get_user(userid) -> UserAccount:
+def kc_get_user_by_id(userid: str) -> UserAccount:
     """
     Get details of specified user account from Keycloak
 
@@ -155,6 +158,29 @@ def kc_get_user(userid) -> UserAccount:
     except (ValueError, AssertionError):
         raise KeycloakError(f"Can't get user info: Unexpected content of response from Keycloak")
 
+def kc_get_user_by_name(username: str) -> Optional[UserAccount]:
+    """
+    Get details of specified user account from Keycloak
+
+    :param username: Keycloak username (not ID)
+    :return UserAccount representation of the user or None is user not found
+    :raise KeycloakError
+    """
+    token = kc_get_token()
+    url = KEYCLOAK_USERS_URL
+    resp = requests.get(url, params={'username': username, 'exact': 'true'}, headers={'Authorization': 'Bearer ' + token}, verify=CA_CERT_FILE)
+    if not resp.ok:
+        raise KeycloakError(f"Can't get user info: ({resp.status_code}) {resp.text[:200]}")
+    print(resp.text)
+    try:
+        users = resp.json()
+        assert isinstance(users, list), ""
+    except (ValueError, AssertionError):
+        raise KeycloakError(f"Can't get user info: Unexpected content of response from Keycloak")
+    if len(users) == 0:
+        raise KeycloakError(f"No user with username '{username}'")
+    return UserAccount.from_keycloak_representation(users[0])
+
 
 def kc_add_user(user: UserAccount) -> None:
     """Add a new user to Keycloak
@@ -203,6 +229,9 @@ def kc_delete_user(userid: str) -> None:
         raise KeycloakError(f"({resp.status_code}) {resp.text[:200]}")
 
 
+# *** NiFi ***
+
+
 # *** Flask endpoints and forms ***
 
 class AddUserForm(FlaskForm):
@@ -224,12 +253,20 @@ def main():
         users = []
     #print(users)
 
+    # Load NiFi users
+    try:
+        nifi_users = nifi_get_users()
+    except NifiError as e:
+        flash(f"ERROR: {e}", "error")
+        nifi_users = []
+    nifi_usernames = set(nu["name"] for nu in nifi_users)
+
     return render_template("main.html", **locals())
 
 
 @app.route("/add_user", methods=["GET", "POST"])
 def add_user():
-    """Add a new user. On GET show new-user page, on POST create new user account."""
+    """Add a new user. On GET show the new-user page, on POST create new user account."""
     form_user = AddUserForm()
     if form_user.validate_on_submit():
         # Form submitted and valid - create user account
@@ -239,18 +276,29 @@ def add_user():
                            lastname=form_user.lastname.data,
                            cn=form_user.cn.data,
                            dn=f"CN={form_user.cn.data}")
+        # Keycloak
         try:
             kc_add_user(user)
-            flash(f'User "{form_user.username.data}" successfully created.', "success")
-            return redirect("/") # Success - go back to main page
-        except KeycloakError as e:
-            flash(f'Error when creating user: {e}', "error")
+            flash(f'User "{form_user.username.data}" successfully created in Keycloak.', "success")
+        except Exception as e:
+            flash(f'Error when creating user in Keycloak: {e}', "error")
+        # NiFi
+        try:
+            nifi_add_user(user)
+            flash(f'User "{form_user.username.data}" successfully created in NiFi.', "success")
+        except NifiUserExistsError:
+            flash(f'User "{user.username}" already exists in NiFi, nothing has changed.', "warning")
+        except Exception as e:
+            flash(f'Error when creating user in NiFi: {e}', "error")
+        return redirect("/") # Success - go back to main page
+
     return render_template("add_edit_user.html", form_user=form_user, user=None)
 
 
-@app.route("/edit_user/<userid>", methods=["GET", "POST"])
-def edit_user(userid: str):
+@app.route("/edit_user/<username>", methods=["GET", "POST"])
+def edit_user(username: str):
     """Edit existing user. On GET show user details, on POST update user params with new values."""
+    keycloak_id = kc_get_user_by_name(username).kcid # TODO catch exception
     if request.method == "POST":
         form_user = AddUserForm() # use data from POST request
         if form_user.validate_on_submit():
@@ -261,17 +309,17 @@ def edit_user(userid: str):
                                lastname=form_user.lastname.data,
                                cn=form_user.cn.data,
                                dn=f"CN={form_user.cn.data}",
-                               kcid=userid)
+                               kcid=keycloak_id)
             try:
                 kc_update_user(user)
                 flash(f'User "{form_user.username.data}" successfully updated.', "success")
                 return redirect("/") # Success - go back to main page
             except KeycloakError as e:
                 flash(f'Error when updating user: {e}', "error")
-        return render_template("add_edit_user.html", form_user=form_user, user={"kcid": userid})
+        return render_template("add_edit_user.html", form_user=form_user, user={"kcid": keycloak_id})
     # else - method="GET"
     try:
-        user = kc_get_user(userid)
+        user = kc_get_user_by_id(keycloak_id)
     except KeycloakError as e:
         flash(f'ERROR: {e}', "error")
         return redirect('/')
@@ -279,14 +327,24 @@ def edit_user(userid: str):
     return render_template("add_edit_user.html", form_user=form_user, user=user)
 
 
-@app.route("/delete_user/<userid>")
-def delete_user(userid: str):
-    """Delete user given by userid and redirect back to main page"""
+@app.route("/delete_user/<username>")
+def delete_user(username: str):
+    """Delete user given by username and redirect back to main page"""
     try:
-        kc_delete_user(userid)
-        flash(f'User successfully deleted.', "success")
+        keycloak_id = kc_get_user_by_name(username).kcid
+        kc_delete_user(keycloak_id)
+        flash('User successfully deleted from KeyCloak.', "success")
     except KeycloakError as e:
-        flash(f'Error when deleting user: {e}', "error")
+        flash(f'Error when deleting user from KeyCloak: {e}', "error")
+
+    try:
+        nifi_delete_user(username)
+        flash(f'User "{username}" successfully deleted from NiFi.', "success")
+    except NifiUserNotFoundError:
+        flash(f'User "{username}" was not found in NiFi, nothing has changed.', "warning")
+    except NifiError as e:
+        flash(f'Error when deleting user from NiFi: {e}', "error")
+
     return redirect("/")
 
 
@@ -294,6 +352,8 @@ def delete_user(userid: str):
 
 # TODO other services (besides Keycloak)
 
+# TODO authentication/authorization to this GUI
+
 # When the script is run directly, run the application on a local development server.
 # Optionally pass two parameters, 'host' (IP to listen on) and 'port',
 # e.g.: ./main.py 0.0.0.0 8080

nifi.py

+"""Functions to manage user accounts in NiFi"""
+
+from typing import List, Dict, Optional
+import subprocess
+import xml.etree.ElementTree as ET
+
+from config import *
+
+# Path to user configuration in NiFi containers
+NIFI_USER_CONFIG_PATH = "/opt/nifi/nifi-current/conf/users.xml"
+
+# Shell command to restart NiFi in the container (simple "supervisorctl restart" doesn't work, since there is another
+# nifi process which supervisord doesn't see and which stops only after some time after the main one; so we need to
+# wait until it stops as well by calling "ps" in a loop)
+#NIFI_RESTART_COMMAND = "supervisorctl stop nifi ; while (ps aux | grep '^nifi' >/dev/null); do sleep 1; done; supervisorctl start nifi"
+NIFI_RESTART_COMMAND = "bin/nifi.sh stop" # stop properly by sending a stop command. It is then restarted automatically by supervisord.
+
+
+class NifiError(Exception):
+    pass
+
+class NifiUserNotFoundError(NifiError):
+    pass
+
+class NifiUserExistsError(NifiError):
+    pass
+
+# For reference, an example NiFi user-config file looks like this:
+"""
+<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
+<tenants>
+    <groups>
+      <group identifier="c78caf19-016f-1000-0000-000000000001" name="NiFi nodes">
+            <user identifier="c78caf19-016f-1000-0001-000000000001"/>
+            <user identifier="c78caf19-016f-1000-0001-000000000002"/>
+            <user identifier="c78caf19-016f-1000-0001-000000000003"/>
+      </group>
+      <group identifier="c78caf19-016f-1000-0000-000000000002" name="Administrators">
+            <user identifier="c78caf19-016f-1000-0002-000000000001"/>
+            <user identifier="c78caf19-016f-1000-0002-000000000002"/>
+      </group>
+    </groups>
+    <users>
+        <user identifier="c78caf19-016f-1000-0001-000000000001" identity="CN=soctools-nifi-1"/>
+        <user identifier="c78caf19-016f-1000-0001-000000000002" identity="CN=soctools-nifi-2"/>
+        <user identifier="c78caf19-016f-1000-0001-000000000003" identity="CN=soctools-nifi-3"/>
+        <user identifier="c78caf19-016f-1000-0002-000000000001" identity="user1"/>
+        <user identifier="c78caf19-016f-1000-0002-000000000002" identity="user2"/>
+    </users>
+</tenants>
+"""
+
+
+def _nifi_xml_get_users(config_xml: str) -> List[Dict]:
+    """Parse the XML file and return list of users.
+
+    Return: list of dicts with keys 'name', 'id', 'group'
+    """
+    config_root = ET.fromstring(config_xml)
+    # Load groups and store group name for each user ID
+    user_to_group = {} # dict: user_id -> group_name
+    for group_elem in config_root.find("groups").findall("group"):
+        g_name = group_elem.get("name")
+        for user_elem in group_elem.findall("user"):
+            user_to_group[user_elem.get("identifier")] = g_name
+    # Load users, find the group for each one
+    users = [] # list of dicts with keys 'name' (identity), 'id' (identifier), 'group'
+    for user_elem in config_root.findall("./users/user"):
+        users.append({
+            "name": user_elem.get("identity"),
+            "id": user_elem.get("identifier"),
+            "group": user_to_group[user_elem.get("identifier")],
+        })
+    return users
+
+
+def _nifi_xml_add_user(config_xml: str, user_name: str, user_group: str) -> str:
+    """Add given user to the XML config file.
+
+    Assumes that "user_group" already exists in the file.
+
+    :return updated xml string
+    """
+    config_root = ET.fromstring(config_xml)
+    # Get info about the given group
+    group_node = config_root.find(f"./groups/group[@name='{user_group}']")
+    group_user_nodes = group_node.findall("user")
+    # Get list of users
+    users_node = config_root.find("./users")
+    # check that there is not a user with the same username
+    if any(u.get("name") == user_name for u in users_node):
+        raise NifiUserExistsError(f"Username '{user_name}' already exists!")
+
+    # Generate new user identifier as the max id of the group plus one
+    #   ids look like: abcd1234-01ab-1000-0002-000000000001 (the last part seems to be the index of the user within the group)
+    user_ids = [u.get("identifier") for u in group_user_nodes]
+    max_id = max(user_ids)
+    group_id,last_user_id = max_id.rsplit("-", 1)
+    new_id = f"{group_id}-{int(last_user_id)+1:012d}"
+
+    # Add a new element the list of users in the group and the list of users
+    group_node.append(ET.Element("user", identifier=new_id))
+    users_node.append(ET.Element("user", identifier=new_id, identity=user_name))
+
+    return ET.tostring(config_root, encoding='utf-8')
+
+def _nifi_xml_delete_user(config_xml: str, user_name: str) -> str:
+    """Remove given user from the XML config file.
+
+    Assumes that "user_group" already exists in the file.
+
+    :return updated xml string
+    """
+    config_root = ET.fromstring(config_xml)
+    # Find user with given name
+    user_node = config_root.find(f"./users/user[@identity='{user_name}']")
+    if user_node is None:
+        raise NifiUserNotFoundError(f"Can't delete user '{user_name}' from NiFi: User with such a name doesn't exist.")
+    # Get user's numerical id
+    identifier = user_node.get('identifier')
+
+    # Remove the user from <users> and from any groups
+    print("identifier:", identifier)
+    print("user_node:", user_node)
+    config_root.find("./users").remove(user_node)
+    for group_node in config_root.findall("./groups/"):
+        print("group_node:", group_node)
+        user_node = group_node.find(f"user[@identifier='{identifier}']")
+        print("  user_node:", user_node)
+        if user_node is not None:
+            group_node.remove(user_node)
+
+    return ET.tostring(config_root, encoding='utf-8')
+
+
+def _nifi_load_user_config(cont_name: str) -> str:
+    """Get the current user-config file from a NiFi docker container"""
+    #print("Getting NiFi config...")
+    result = subprocess.run(["docker", "exec", cont_name, "cat", NIFI_USER_CONFIG_PATH],
+                            stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    if result.returncode == 0:
+        config_xml = result.stdout.decode('utf-8')
+        #print(config_xml)
+        return config_xml
+    else:
+        raise NifiError(f'Error when trying to get the current config of NiFi users from container "{cont_name}": {result.stderr.decode()}')
+
+
+def _nifi_write_user_config(new_xml: str):
+    """Write given XML string to the user-config file in all NiFi containers, restart NiFi in all containers."""
+    # Write new file into all containers
+    for i,cont_name in enumerate(NIFI_CONTAINERS):
+        # Run a command to write the new user-config file inside the NiFi docker container
+        # The file contents are passed via stdin ("docker exec -i" must be used for stdin to work), "tee" is used to
+        # store it to a file (stdout is ignored)
+        print(f'Writing new NiFi user config in container "{cont_name}"')
+        result = subprocess.run(["docker", "exec", "-i", cont_name, "tee", NIFI_USER_CONFIG_PATH],
+                                input=new_xml, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
+        if result.returncode != 0:
+            raise NifiError(f'Error when trying to write the updated list of NiFi users into container "{cont_name}": {result.stderr.decode()}')
+    # Restart all containers (may take a long time!)
+    for i,cont_name in enumerate(NIFI_CONTAINERS):
+        print(f'Restarting NiFi in container "{cont_name}"')
+        result = subprocess.run(["docker", "exec", cont_name, "bash", "-c", NIFI_RESTART_COMMAND],
+                                stderr=subprocess.PIPE)
+        if result.returncode != 0:
+            raise NifiError(f'Error when trying to restart NiFi after config update in container "{cont_name}": {result.stderr.decode()}')
+
+
+def nifi_get_users() -> List[Dict]:
+    """
+    List users defined in NiFi
+    """
+    prev_users = None
+    for i,cont_name in enumerate(NIFI_CONTAINERS):
+        # Get the current user-config file
+        config_xml = _nifi_load_user_config(cont_name)
+        if config_xml is False:
+            return []
+
+        # Parse the list of users from the config file
+        try:
+            users = _nifi_xml_get_users(config_xml)
+        except Exception as e:
+            raise NifiError(f'Can\'t parse NiFi user config file ("{NIFI_USER_CONFIG_PATH}" in container "{cont_name}"): {e}')
+
+        # Remove "internal" users from the "NiFi nodes" group
+        users = [u for u in users if u["group"] != "NiFi nodes"]
+        # TODO this way, the comparison forces the order of users to be the same, which is not needed (although
+        #  it is always the same in normal situation)
+        # Check that the list is the same as in the previous container (all should be the same)
+        if prev_users is not None and users != prev_users:
+            raise NifiError('Error when trying to get the list of NiFi users: The lists of users differ in at least '
+                            f'two NiFi nodes ({NIFI_CONTAINERS[i-1]},{NIFI_CONTAINERS[i]}). Check the file '
+                            f'"{NIFI_USER_CONFIG_PATH}" in each NiFi container, they all must be the same.')
+        prev_users = users
+    return prev_users
+
+
+def nifi_add_user(user: 'UserAccount'):
+    """Add a new user to NiFi
+
+    Read user config from a NiFi container (the first one), add a new user it, and write it in all containers (overwrite!).
+
+    :raises NifiError
+    """
+    user_name = user.username
+    user_group = "Administrators" # no support for other groups in NiFi, yet
+
+    # Get the current user-config file (use the first container, the content should be the same in all)
+    cont_name = NIFI_CONTAINERS[0]
+    config_xml = _nifi_load_user_config(cont_name)
+
+    #Add new user to the XML
+    try:
+        new_xml = _nifi_xml_add_user(config_xml, user_name, user_group)
+    except NifiUserExistsError:
+        raise
+    except Exception as e:
+        raise NifiError(f'Can\'t add user to the config file ("{NIFI_USER_CONFIG_PATH}" in container "{cont_name}"): {e}')
+
+    #print("XML with added user:")
+    #print(new_xml)
+
+    # Write the updated config and restart NiFi nodes
+    _nifi_write_user_config(new_xml)
+
+
+def nifi_delete_user(user_name: str):
+    """Add a new user to NiFi
+
+    Read user config from a NiFi container (the first one), remove the specified user, and write new version in all containers (overwrite!).
+
+    :raises NifiError
+    """
+    # Get the current user-config file (use the first container, the content should be the same in all)
+    cont_name = NIFI_CONTAINERS[0]
+    config_xml = _nifi_load_user_config(cont_name)
+
+    # Remove user from the XML
+    try:
+        new_xml = _nifi_xml_delete_user(config_xml, user_name)
+    except NifiUserNotFoundError:
+        raise
+    except Exception as e:
+        raise NifiError(f'Can\'t remove user from config file ("{NIFI_USER_CONFIG_PATH}" in container "{cont_name}"): {e}')
+
+    #print("XML with deleted user:")
+    #print(new_xml)
+
+    # Write the updated config and restart NiFi nodes
+    _nifi_write_user_config(new_xml)
+

static/style.css

@@ -18,6 +18,10 @@ p {
   padding: 0.5em;
 }
 
+img.icon {
+  vertical-align: center;
+}
+
 .errors {
   background-color: #fcc;
   color: #c00;
@@ -29,6 +33,9 @@ ul.flashes {
 li.flash-error {
   color: #900;
 }
+li.flash-warning {
+  color: #990;
+}
 li.flash-success {
   color: #090;
 }

templates/add_edit_user.html

@@ -10,7 +10,7 @@
 <h2>Add new user</h2>
 {% endif %}
 
-<form action="{{ url_for("edit_user", userid=user.kcid) if user else url_for("add_user") }}" method="POST">
+<form action="{{ url_for("edit_user", username=user.username) if user else url_for("add_user") }}" method="POST">
 {% if form_user.errors %}
     <ul class="errors">
     {% for field, errors in form_user.errors.items() %}

templates/base.html

+{%- macro icon(type, size='24px') -%}
+<img src="{{ url_for("static", filename="icons/"+type+".svg") }}" class="icon" style="width: {{size}}; height: {{size}}">
+{%- endmacro -%}
 <!doctype html>
 <html>
 <head>

templates/main.html

 {% extends "base.html" %}
 {% block body %}
 
-<p><a href="{{ url_for("add_user") }}">Add new user ...</a></p>
+<p><a href="{{ url_for("add_user") }}">{{ icon('plus-circle', size="1em") }} Add new user ...</a></p>
 
 <table>
-<tr><th>Username</th><th>First name</th><th>Last name</th><th>email</th><th>CN</th><th>DN</th><th>Time created (UTC)</th><th></th>
+<tr><th>Username</th><th>First name</th><th>Last name</th><th>email</th><th>CN</th><th>DN</th><th>Time created (UTC)</th><th>NiFi</th><th></th>
 
 {% for user in users %}
 <tr>
@@ -15,14 +15,30 @@
 <td>{{ user.cn }}</td>
 <td>{{ user.dn }}</td>
 <td>{{ user.ts_created.isoformat() }}</td>
+<td>{{ icon('check' if user.username in nifi_usernames else 'close') }}</td>
 <td>
-<a href="{{ url_for('edit_user', userid=user.kcid) }}" title="Edit user">&#128393;</a>
-<a href="{{ url_for('delete_user', userid=user.kcid) }}" title="Delete user"
- onclick="return confirm('Are you sure you want to permanently delete user account &quot;{{user.username}}&quot; ({{user.cn}}, {{user.email}})?')">&#128465;</a>
+<a href="{{ url_for('edit_user', username=user.username) }}" title="Edit user">{{ icon('pencil') }}</a>
+<a href="{{ url_for('delete_user', username=user.username) }}" title="Delete user"
+ onclick="return confirm('Are you sure you want to permanently delete user account &quot;{{user.username}}&quot; ({{user.cn}}, {{user.email}})?')">{{ icon('trash') }}</a>
 </td>
 </tr>
 {#<tr><td colspan=8>{{ user }}</td></tr>#}
 {% endfor %}
 </table>
 
+<h2>Users in individual services</h2>
+
+<h3>NiFi</h3>
+<table>
+<tr><th>Username</th><th>Group</th><th>ID</th>
+
+{% for user in nifi_users %}
+<tr>
+<td>{{ user.name }}</td>
+<td>{{ user.group }}</td>
+<td>{{ user.id }}</td>
+</tr>
+{% endfor %}
+</table>
+
 {% endblock %}
\ No newline at end of file