Add transform_function (#19)

nornir-automation · Oct 27, 2021 · 3c2897d · 3c2897d
1 parent 2ec7939
commit 3c2897d
Show file tree

Hide file tree

Showing 3 changed files with 344 additions and 0 deletions.
diff --git a/docs/source/tutorials/load_credentials.ipynb b/docs/source/tutorials/load_credentials.ipynb
@@ -0,0 +1,318 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "# load_credentials\n",
+    "\n",
+    "`load_credentials` is an transform_function to add credentials to every host. The transform_function_options `username` and `password` are supportet as well as the environment variables `NORNIR_USERNAME` and `NORNIR_PASSWORD`. transform_function_options have priority.\n",
+    "\n",
+    "## Example"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Loading inventory without load_credentials"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dev1.group_1\tUser:a_user\tPassword:a_password\n",
+      "dev2.group_1\tUser:None\tPassword:from_group1\n",
+      "dev3.group_2\tUser:None\tPassword:None\n",
+      "dev4.group_2\tUser:None\tPassword:from_parent_group\n",
+      "dev5.no_group\tUser:None\tPassword:None\n"
+     ]
+    }
+   ],
+   "source": [
+    "from nornir import InitNornir\n",
+    "\n",
+    "nr = InitNornir(\n",
+    "    inventory={\n",
+    "        \"plugin\": \"SimpleInventory\",\n",
+    "        \"options\": {\"host_file\": \"data/hosts.yaml\", \"group_file\": \"data/groups.yaml\"}\n",
+    "    }\n",
+    ")\n",
+    "\n",
+    "for host in nr.inventory.hosts.values():\n",
+    "    print(f\"{host.name}\\tUser:{host.username}\\tPassword:{host.password}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Using transform_funtion_options\n",
+    "\n",
+    "Set the password for all hosts using the parameter `password`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dev1.group_1\tUser:a_user\tPassword:TopSecretP4ss\n",
+      "dev2.group_1\tUser:None\tPassword:TopSecretP4ss\n",
+      "dev3.group_2\tUser:None\tPassword:TopSecretP4ss\n",
+      "dev4.group_2\tUser:None\tPassword:TopSecretP4ss\n",
+      "dev5.no_group\tUser:None\tPassword:TopSecretP4ss\n"
+     ]
+    }
+   ],
+   "source": [
+    "from nornir import InitNornir\n",
+    "\n",
+    "nr = InitNornir(\n",
+    "    inventory={\n",
+    "        \"plugin\": \"SimpleInventory\",\n",
+    "        \"options\": {\"host_file\": \"data/hosts.yaml\", \"group_file\": \"data/groups.yaml\"},\n",
+    "        \"transform_function\": \"load_credentials\",\n",
+    "        \"transform_function_options\": {\"password\": \"TopSecretP4ss\"}\n",
+    "    }\n",
+    ")\n",
+    "\n",
+    "for host in nr.inventory.hosts.values():\n",
+    "    print(f\"{host.name}\\tUser:{host.username}\\tPassword:{host.password}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setting username and password"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dev1.group_1\tUser:myUser\tPassword:TopSecretP4ss\n",
+      "dev2.group_1\tUser:myUser\tPassword:TopSecretP4ss\n",
+      "dev3.group_2\tUser:myUser\tPassword:TopSecretP4ss\n",
+      "dev4.group_2\tUser:myUser\tPassword:TopSecretP4ss\n",
+      "dev5.no_group\tUser:myUser\tPassword:TopSecretP4ss\n"
+     ]
+    }
+   ],
+   "source": [
+    "from nornir import InitNornir\n",
+    "\n",
+    "nr = InitNornir(\n",
+    "    inventory={\n",
+    "        \"plugin\": \"SimpleInventory\",\n",
+    "        \"options\": {\"host_file\": \"data/hosts.yaml\", \"group_file\": \"data/groups.yaml\"},\n",
+    "        \"transform_function\": \"load_credentials\",\n",
+    "        \"transform_function_options\": {\"username\": \"myUser\", \"password\": \"TopSecretP4ss\"}\n",
+    "    }\n",
+    ")\n",
+    "\n",
+    "for host in nr.inventory.hosts.values():\n",
+    "    print(f\"{host.name}\\tUser:{host.username}\\tPassword:{host.password}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Using environment variables\n",
+    "\n",
+    "Set the variable `NORNIR_PASSWORD`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "env: NORNIR_PASSWORD=AnotherPassWord!\n"
+     ]
+    }
+   ],
+   "source": [
+    "%set_env NORNIR_PASSWORD=AnotherPassWord!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dev1.group_1\tUser:a_user\tPassword:AnotherPassWord!\n",
+      "dev2.group_1\tUser:None\tPassword:AnotherPassWord!\n",
+      "dev3.group_2\tUser:None\tPassword:AnotherPassWord!\n",
+      "dev4.group_2\tUser:None\tPassword:AnotherPassWord!\n",
+      "dev5.no_group\tUser:None\tPassword:AnotherPassWord!\n"
+     ]
+    }
+   ],
+   "source": [
+    "from nornir import InitNornir\n",
+    "\n",
+    "nr = InitNornir(\n",
+    "    inventory={\n",
+    "        \"plugin\": \"SimpleInventory\",\n",
+    "        \"options\": {\"host_file\": \"data/hosts.yaml\", \"group_file\": \"data/groups.yaml\"},\n",
+    "        \"transform_function\": \"load_credentials\"\n",
+    "    }\n",
+    ")\n",
+    "\n",
+    "for host in nr.inventory.hosts.values():\n",
+    "    print(f\"{host.name}\\tUser:{host.username}\\tPassword:{host.password}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Also set the environment variable `NORNIR_USERNAME`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "env: NORNIR_USERNAME=AnotherUser\n"
+     ]
+    }
+   ],
+   "source": [
+    "%set_env NORNIR_USERNAME=AnotherUser"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dev1.group_1\tUser:AnotherUser\tPassword:AnotherPassWord!\n",
+      "dev2.group_1\tUser:AnotherUser\tPassword:AnotherPassWord!\n",
+      "dev3.group_2\tUser:AnotherUser\tPassword:AnotherPassWord!\n",
+      "dev4.group_2\tUser:AnotherUser\tPassword:AnotherPassWord!\n",
+      "dev5.no_group\tUser:AnotherUser\tPassword:AnotherPassWord!\n"
+     ]
+    }
+   ],
+   "source": [
+    "from nornir import InitNornir\n",
+    "\n",
+    "nr = InitNornir(\n",
+    "    inventory={\n",
+    "        \"plugin\": \"SimpleInventory\",\n",
+    "        \"options\": {\"host_file\": \"data/hosts.yaml\", \"group_file\": \"data/groups.yaml\"},\n",
+    "        \"transform_function\": \"load_credentials\"\n",
+    "    }\n",
+    ")\n",
+    "\n",
+    "for host in nr.inventory.hosts.values():\n",
+    "    print(f\"{host.name}\\tUser:{host.username}\\tPassword:{host.password}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The transform_function_option can overwrite the environment variables"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dev1.group_1\tUser:AnotherUser\tPassword:TopSecretP4ss\n",
+      "dev2.group_1\tUser:AnotherUser\tPassword:TopSecretP4ss\n",
+      "dev3.group_2\tUser:AnotherUser\tPassword:TopSecretP4ss\n",
+      "dev4.group_2\tUser:AnotherUser\tPassword:TopSecretP4ss\n",
+      "dev5.no_group\tUser:AnotherUser\tPassword:TopSecretP4ss\n"
+     ]
+    }
+   ],
+   "source": [
+    "from nornir import InitNornir\n",
+    "\n",
+    "nr = InitNornir(\n",
+    "    inventory={\n",
+    "        \"plugin\": \"SimpleInventory\",\n",
+    "        \"options\": {\"host_file\": \"data/hosts.yaml\", \"group_file\": \"data/groups.yaml\"},\n",
+    "        \"transform_function\": \"load_credentials\",\n",
+    "        \"transform_function_options\": {\"password\": \"TopSecretP4ss\"}\n",
+    "    }\n",
+    ")\n",
+    "\n",
+    "for host in nr.inventory.hosts.values():\n",
+    "    print(f\"{host.name}\\tUser:{host.username}\\tPassword:{host.password}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "9e96c052789e49a59d2331cc046e176b6533b6223242088d74e01af11b5e017c"
+  },
+  "kernelspec": {
+   "display_name": "Python 3.8.10 64-bit",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/nornir_utils/plugins/inventory/transform_functions.py b/nornir_utils/plugins/inventory/transform_functions.py
@@ -0,0 +1,23 @@
+import os
+from typing import Optional
+from nornir.core.inventory import Host
+
+
+def load_credentials(
+    host: Host, username: Optional[str] = None, password: Optional[str] = None
+) -> None:
+    """
+    load_credentials is an transform_functions to add credentials to every host.
+    Environment variables `NORNIR_USERNAME` and `NORNIR_PASSWORD` or arguments can be used.
+
+    Args:
+
+        username: Device username
+        password: Device password
+    """
+    username = username if username is not None else os.getenv("NORNIR_USERNAME")
+    if username is not None:
+        host.username = username
+    password = password if password is not None else os.getenv("NORNIR_PASSWORD")
+    if password is not None:
+        host.password = password
diff --git a/pyproject.toml b/pyproject.toml
@@ -8,6 +8,9 @@ license = "Apache-2.0"
 [tool.poetry.plugins."nornir.plugins.inventory"]
 "YAMLInventory" = "nornir_utils.plugins.inventory.yaml_inventory:YAMLInventory"
 
+[tool.poetry.plugins."nornir.plugins.transform_function"]
+"load_credentials" = "nornir_utils.plugins.inventory.transform_functions:load_credentials"
+
 [tool.poetry.dependencies]
 python = "^3.6"
 nornir = { version = "~3", allow-prereleases = true }