From a6faba3d1dda42227cd47ba5a5af18cc4ed66b4c Mon Sep 17 00:00:00 2001 From: Wim Van Deun <7521270+enzzzy@users.noreply.github.com> Date: Sun, 17 Feb 2019 21:23:27 +0100 Subject: [PATCH] Add gitlab file plugin --- nornir/plugins/tasks/files/__init__.py | 3 +- nornir/plugins/tasks/files/gitlab.py | 182 ++++++++++++ tests/plugins/tasks/files/test_gitlab.py | 346 +++++++++++++++++++++++ 3 files changed, 530 insertions(+), 1 deletion(-) create mode 100644 nornir/plugins/tasks/files/gitlab.py create mode 100644 tests/plugins/tasks/files/test_gitlab.py diff --git a/nornir/plugins/tasks/files/__init__.py b/nornir/plugins/tasks/files/__init__.py index 2b8533f2..6c2c94f0 100644 --- a/nornir/plugins/tasks/files/__init__.py +++ b/nornir/plugins/tasks/files/__init__.py @@ -1,5 +1,6 @@ from .sftp import sftp from .write_file import write_file +from .gitlab import gitlab -__all__ = ("sftp", "write_file") +__all__ = ("sftp", "write_file", "gitlab") diff --git a/nornir/plugins/tasks/files/gitlab.py b/nornir/plugins/tasks/files/gitlab.py new file mode 100644 index 00000000..0e102bca --- /dev/null +++ b/nornir/plugins/tasks/files/gitlab.py @@ -0,0 +1,182 @@ +import base64 +import difflib +import threading +from typing import Tuple + +from nornir.core.task import Result, Task + +import requests + + +LOCK = threading.Lock() + + +def _generate_diff(original: str, fromfile: str, tofile: str, content: str) -> str: + diff = difflib.unified_diff( + original.splitlines(), content.splitlines(), fromfile=fromfile, tofile=tofile + ) + return "\n".join(diff) + + +def _get_repository(session: requests.Session, url: str, repository: str) -> int: + resp = session.get(f"{url}/api/v4/projects?search={repository}") + if resp.status_code != 200: + raise RuntimeError(f"Unexpected Gitlab status code {resp.status_code}") + + pid = 0 + found = False + respjson = resp.json() + if not len(respjson): + raise RuntimeError("Gitlab repository not found") + + for p in respjson: + if p.get("name", "") == repository: + found = True + pid = p.get("id", 0) + + if not pid or not found: + raise RuntimeError("Gitlab repository not found") + + return pid + + +def _remote_exists( + task: Task, + session: requests.Session, + url: str, + pid: int, + filename: str, + branch: str, +) -> Tuple[bool, str]: + resp = session.get( + f"{url}/api/v4/projects/{pid}/repository/files/{filename}?ref={branch}" + ) + if resp.status_code == 200: + return ( + True, + base64.decodebytes(resp.json()["content"].encode("ascii")).decode(), + ) + return (False, "") + + +def _create( + task: Task, + session: requests.Session, + url: str, + pid: int, + filename: str, + content: str, + branch: str, + commit_message: str, + dry_run: bool, +) -> str: + if dry_run: + return _generate_diff("", "", filename, content) + + with LOCK: + url = f"{url}/api/v4/projects/{pid}/repository/files/{filename}" + data = {"branch": branch, "content": content, "commit_message": commit_message} + resp = session.post(url, data=data) + + if resp.status_code != 201: + raise RuntimeError(f"Unable to create file: {filename}!") + return _generate_diff("", "", filename, content) + + +def _update( + task: Task, + session: requests.Session, + url: str, + pid: int, + filename: str, + content: str, + branch: str, + commit_message: str, + dry_run: bool, +) -> str: + exists, original = _remote_exists(task, session, url, pid, filename, branch) + + if not exists: + raise RuntimeError(f"File '{filename}' does not exist!") + + if dry_run: + return _generate_diff(original, filename, filename, content) + + if original != content: + with LOCK: + url = f"{url}/api/v4/projects/{pid}/repository/files/{filename}" + data = { + "branch": branch, + "content": content, + "commit_message": commit_message, + } + resp = session.put(url=url, data=data) + if resp.status_code != 200: + print(f"{resp.status_code} : {resp.text}") + raise RuntimeError(f"Unable to update file: {filename}") + return _generate_diff(original, filename, filename, content) + + +def gitlab( + task: Task, + url: str, + token: str, + repository: str, + filename: str, + content: str, + action: str = "create", + dry_run: bool = False, + branch: str = "master", + commit_message: str = "", +) -> Result: + """ + Writes contents to a new file or update contents of an existing file in a + gitlab repository. + + Example: + + nornir.run(files.gitlab, + action="create", + url="https://gitlab.localhost.com", + token="ABCD1234", + repository="test", + filename="config", + branch="master") + + Arguments: + dry_run: Whether to apply changes or not + url: Gitlab instance URL + token: Personal access token + repository: destination repository + filename: destination file name + content: content to write + action: ``create``, ``update`` + update: Update the file if it already exists, + when create action is used. + branch: destination branch + + Returns: + Result object with the following attributes set: + * changed (``bool``): + * diff (``str``): unified diff + + """ + dry_run = task.is_dry_run(dry_run) + + session = requests.session() + session.headers.update({"PRIVATE-TOKEN": token}) + + if commit_message == "": + commit_message = "File created with nornir" + + pid = _get_repository(session, url, repository) + + if action == "create": + diff = _create( + task, session, url, pid, filename, content, branch, commit_message, dry_run + ) + elif action == "update": + diff = _update( + task, session, url, pid, filename, content, branch, commit_message, dry_run + ) + return Result(host=task.host, diff=diff, changed=bool(diff)) diff --git a/tests/plugins/tasks/files/test_gitlab.py b/tests/plugins/tasks/files/test_gitlab.py new file mode 100644 index 00000000..f00f0432 --- /dev/null +++ b/tests/plugins/tasks/files/test_gitlab.py @@ -0,0 +1,346 @@ +import os + +from nornir.plugins.tasks import files + + +BASE_PATH = os.path.join(os.path.dirname(__file__), "gitlab") + +diff_create = """--- + ++++ dummy + +@@ -0,0 +1 @@ + ++dummy""" # noqa + +diff_update = """--- dummy + ++++ dummy + +@@ -1 +1 @@ + +-dummy ++new line""" + + +def create_file( + nornir, + requests_mock, + url, + repository, + pid, + branch, + filename, + content, + action, + dry_run, + commit_message, + status_code, + project_status_code, + project_resp, + resp, +): + token = "dummy" + + repo_url = f"{url}/api/v4/projects?search={repository}" + requests_mock.get(repo_url, status_code=project_status_code, json=project_resp) + + create_file_url = f"{url}/api/v4/projects/{pid}/repository/files/{filename}" + requests_mock.post(create_file_url, status_code=status_code, json=resp) + + res = nornir.run( + files.gitlab, + url=url, + token=token, + repository=repository, + filename=filename, + content=content, + action="create", + dry_run=dry_run, + branch=branch, + commit_message=commit_message, + ) + return res + + +def update_file( + nornir, + requests_mock, + url, + repository, + pid, + branch, + filename, + content, + dry_run, + commit_message, + status_code, + project_status_code, + exists_status_code, + project_resp, + exists_resp, + resp, +): + token = "dummy" + + repo_url = f"{url}/api/v4/projects?search={repository}" + requests_mock.get(repo_url, status_code=project_status_code, json=project_resp) + + exists_file_url = ( + f"{url}/api/v4/projects/{pid}/repository/files/{filename}?ref={branch}" + ) + requests_mock.get(exists_file_url, status_code=exists_status_code, json=exists_resp) + + update_file_url = f"{url}/api/v4/projects/{pid}/repository/files/{filename}" + requests_mock.put(update_file_url, status_code=status_code, json=resp) + + res = nornir.run( + files.gitlab, + url=url, + token=token, + repository=repository, + filename=filename, + content=content, + action="update", + dry_run=dry_run, + branch=branch, + commit_message=commit_message, + ) + return res + + +class Test(object): + def test_gitlab_create_dry_run(self, nornir, requests_mock): + nornir = nornir.filter(name="dev1.group_1") + res = create_file( + nornir=nornir, + requests_mock=requests_mock, + url="http://localhost", + repository="test", + pid=1, + branch="master", + filename="dummy", + content="dummy", + action="create", + dry_run=True, + commit_message="commit", + status_code=201, + project_status_code=200, + project_resp=[{"name": "test", "id": 1}], + resp={"branch": "master", "file_path": "dummy"}, + ) + + assert not res["dev1.group_1"][0].failed + assert res["dev1.group_1"][0].changed + assert res["dev1.group_1"][0].diff == diff_create + + def test_gitlab_create(self, nornir, requests_mock): + nornir = nornir.filter(name="dev1.group_1") + res = create_file( + nornir=nornir, + requests_mock=requests_mock, + url="http://localhost", + repository="test", + pid=1, + branch="master", + filename="dummy", + content="dummy", + action="create", + dry_run=False, + commit_message="commit", + status_code=201, + project_status_code=200, + project_resp=[{"name": "test", "id": 1}], + resp={"branch": "master", "file_path": "dummy"}, + ) + + assert not res["dev1.group_1"][0].failed + assert res["dev1.group_1"][0].changed + assert res["dev1.group_1"][0].diff == diff_create + + def test_gitlab_create_file_exists(self, nornir, requests_mock): + nornir = nornir.filter(name="dev1.group_1") + res = create_file( + nornir=nornir, + requests_mock=requests_mock, + url="http://localhost", + repository="test", + pid=1, + branch="master", + filename="dummy", + content="dummy", + action="create", + dry_run=False, + commit_message="commit", + status_code=400, + project_status_code=200, + project_resp=[{"name": "test", "id": 1}], + resp={"branch": "master", "file_path": "dummy"}, + ) + + assert res["dev1.group_1"][0].failed + assert not res["dev1.group_1"][0].changed + + def test_gitlab_create_invalid_project(self, nornir, requests_mock): + nornir = nornir.filter(name="dev1.group_1") + res = create_file( + nornir=nornir, + requests_mock=requests_mock, + url="http://localhost", + repository="test", + pid=1, + branch="master", + filename="dummy", + content="dummy", + action="create", + dry_run=False, + commit_message="commit", + status_code=201, + project_status_code=200, + project_resp=[{"name": "aaa", "id": 1}], + resp={"branch": "master", "file_path": "dummy"}, + ) + + assert res["dev1.group_1"][0].failed + assert not res["dev1.group_1"][0].changed + + def test_gitlab_create_invalid_branch(self, nornir, requests_mock): + nornir = nornir.filter(name="dev1.group_1") + res = create_file( + nornir=nornir, + requests_mock=requests_mock, + url="http://localhost", + repository="test", + pid=1, + branch="bar", + filename="dummy", + content="dummy", + action="create", + dry_run=False, + commit_message="commit", + status_code=400, + project_status_code=200, + project_resp=[{"name": "test", "id": 1}], + resp={"branch": "master", "file_path": "dummy"}, + ) + + assert res["dev1.group_1"][0].failed + assert not res["dev1.group_1"][0].changed + + def test_gitlab_update_dry_run(self, nornir, requests_mock): + nornir = nornir.filter(name="dev1.group_1") + res = update_file( + nornir=nornir, + requests_mock=requests_mock, + url="http://localhost", + repository="test", + pid=1, + branch="master", + filename="dummy", + content="new line", + dry_run=True, + commit_message="commit", + status_code=200, + project_status_code=200, + exists_status_code=200, + project_resp=[{"name": "test", "id": 1}], + exists_resp={"content": "ZHVtbXk=\n"}, + resp={"branch": "master", "file_path": "dummy"}, + ) + + assert not res["dev1.group_1"][0].failed + assert res["dev1.group_1"][0].changed + assert res["dev1.group_1"][0].diff == diff_update + + def test_gitlab_update(self, nornir, requests_mock): + nornir = nornir.filter(name="dev1.group_1") + res = update_file( + nornir=nornir, + requests_mock=requests_mock, + url="http://localhost", + repository="test", + pid=1, + branch="master", + filename="dummy", + content="new line", + dry_run=False, + commit_message="commit", + status_code=200, + project_status_code=200, + exists_status_code=200, + project_resp=[{"name": "test", "id": 1}], + exists_resp={"content": "ZHVtbXk=\n"}, + resp={"branch": "master", "file_path": "dummy"}, + ) + assert not res["dev1.group_1"][0].failed + assert res["dev1.group_1"][0].changed + assert res["dev1.group_1"][0].diff == diff_update + + def test_gitlab_update_invalid_project(self, nornir, requests_mock): + nornir = nornir.filter(name="dev1.group_1") + res = update_file( + nornir=nornir, + requests_mock=requests_mock, + url="http://localhost", + repository="test", + pid=1, + branch="master", + filename="dummy", + content="new line", + dry_run=False, + commit_message="commit", + status_code=200, + project_status_code=200, + exists_status_code=200, + project_resp=[{"name": "123", "id": 1}], + exists_resp={"content": "ZHVtbXk=\n"}, + resp={"branch": "master", "file_path": "dummy"}, + ) + assert res["dev1.group_1"][0].failed + assert not res["dev1.group_1"][0].changed + + def test_gitlab_update_invalid_branch(self, nornir, requests_mock): + nornir = nornir.filter(name="dev1.group_1") + res = update_file( + nornir=nornir, + requests_mock=requests_mock, + url="http://localhost", + repository="test", + pid=1, + branch="bar", + filename="dummy", + content="new line", + dry_run=False, + commit_message="commit", + status_code=200, + project_status_code=200, + exists_status_code=400, + project_resp=[{"name": "test", "id": 1}], + exists_resp="", + resp={"branch": "master", "file_path": "dummy"}, + ) + assert res["dev1.group_1"][0].failed + assert not res["dev1.group_1"][0].changed + + def test_gitlab_update_invalid_file(self, nornir, requests_mock): + nornir = nornir.filter(name="dev1.group_1") + res = update_file( + nornir=nornir, + requests_mock=requests_mock, + url="http://localhost", + repository="test", + pid=1, + branch="master", + filename="bar", + content="new line", + dry_run=False, + commit_message="commit", + status_code=200, + project_status_code=200, + exists_status_code=400, + project_resp=[{"name": "test", "id": 1}], + exists_resp={"content": "ZHVtbXk=\n"}, + resp={"branch": "master", "file_path": "dummy"}, + ) + assert res["dev1.group_1"][0].failed + assert not res["dev1.group_1"][0].changed