knotctl/knotctl

#!/usr/bin/env python3

import argparse
import getpass
import json
import sys
import urllib.parse
from collections.abc import Sequence
from os import environ, mkdir
from os.path import isdir, isfile, join
from typing import Union

import argcomplete
import requests
import yaml
from requests.models import HTTPBasicAuth


def error(description: str, error: str) -> Sequence[dict]:
    response = []
    reply = {}
    # https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/406
    reply["Code"] = 406
    reply["Description"] = description
    reply["Error"] = error
    response.append(reply)
    return response


def get_config(config_filename: str):
    with open(config_filename, "r") as fh:
        return yaml.safe_load(fh.read())


def output(response: Sequence[dict], jsonout: bool = False):
    try:
        if jsonout:
            print(json.dumps(response))
        else:
            for entry in response:
                print(nested_out(entry))
    except BrokenPipeError:
        pass


def nested_out(input, tabs="") -> str:
    string = ""
    if isinstance(input, str):
        string += "{}\n".format(input)
    elif isinstance(input, dict):
        for key, value in input.items():
            string += "{}{}: {}".format(tabs, key, nested_out(value, tabs + " "))
    elif isinstance(input, list):
        for entry in input:
            string += "{}\n{}".format(tabs, nested_out(entry, tabs + " "))
    return string


# Define the parser for each command
def run_add(url: str, jsonout: bool, headers: dict):
    response = requests.put(url, headers=headers)
    output(response.json(), jsonout)


def run_config(
    config_filename: str,
    baseurl: Union[None, str] = None,
    username: Union[None, str] = None,
    password: Union[None, str] = None,
):
    config = {"baseurl": baseurl, "username": username, "password": password}
    needed = []
    if not baseurl:
        needed.append("baseurl")
    if not username:
        needed.append("username")
    for need in needed:
        if need == "":
            output(
                error("Can not configure without {}".format(need), "No {}".format(need))
            )
            sys.exit(1)
        config[need] = input("Enter {}:".format(need))

    if not password:
        try:
            config["password"] = getpass.getpass()
        except EOFError:
            output(error("Can not configure without password", "No password"))
            sys.exit(1)

    with open(config_filename, "w") as fh:
        fh.write(yaml.dump(config))


def run_delete(url: str, jsonout: bool, headers: dict):
    response = requests.delete(url, headers=headers)
    output(response.json(), jsonout)


def run_list(url: str, jsonout: bool, headers: dict):
    response = requests.get(url, headers=headers)
    output(response.json(), jsonout)


def run_update(url: str, jsonout: bool, headers: dict):
    response = requests.patch(url, headers=headers)
    output(response.json(), jsonout)


# Set up the url
def setup_url(
    baseurl: str,
    data: Union[None, Sequence[str]],
    name: Union[None, str],
    rtype: Union[None, str],
    zone: Union[None, str],
) -> str:
    url = baseurl + "/zones"
    if zone:
        url += "/{}".format(zone)
    if name and zone:
        url += "/records/{}".format(name)
    if zone and name and rtype:
        url += "/{}".format(rtype)
    if data and zone and name and rtype:
        url += "?"
        for arg in data:
            if not url.endswith("?"):
                url += "&"
            url += urllib.parse.quote_plus(arg)
    if rtype and (not name or not zone):
        output(error("rtype only makes sense with name and zone", "Missing parameter"))
        sys.exit(1)
    if name and not zone:
        output(error("name only makes sense with a zone", "Missing parameter"))
        sys.exit(1)
    return url


def main() -> int:
    config_basepath = join(environ["HOME"], ".knot")
    config_filename = join(config_basepath, "config")

    if not isdir(config_basepath):
        mkdir(config_basepath)

    if not isfile(config_filename):
        print("You need to configure knotctl before proceeding")
        run_config(config_filename)

    config = get_config(config_filename)
    baseurl = config["baseurl"]
    username = config["username"]
    password = config["password"]

    # Authenticate
    basic = HTTPBasicAuth(username, password)
    response = requests.get(baseurl + "/user/login", auth=basic)
    try:
        token = response.json()["token"]
    except KeyError:
        output(response.json())
        return 1
    headers = {"Authorization": "Bearer {}".format(token)}

    # Grab user input
    parser = argparse.ArgumentParser()
    parser.add_argument("--json", action=argparse.BooleanOptionalAction)
    subparsers = parser.add_subparsers(dest="command")
    for sub in ["add", "delete"]:
        subparser = subparsers.add_parser(sub)
        subparser.add_argument(
            "-d",
            "--data",
            nargs="*",
            help="Specify any number of key - value pairs: name=dns1.example.com.",
        )
        subparser.add_argument("-n", "--name")
        subparser.add_argument("-r", "--rtype")
        subparser.add_argument("-z", "--zone", required=True)
    list = subparsers.add_parser("list")
    list.add_argument(
        "-d",
        "--data",
        nargs="*",
        help="Specify any number of key - value pairs: name=dns1.example.com.",
    )
    list.add_argument("-n", "--name")
    list.add_argument("-r", "--rtype")
    list.add_argument("-z", "--zone")

    config = subparsers.add_parser("config")
    config.add_argument("-b", "--baseurl")
    config.add_argument("-p", "--password")
    config.add_argument("-u", "--username")

    update = subparsers.add_parser("update")
    update.add_argument(
        "-d",
        "--data",
        nargs="*",
        help="Specify any number of key - value pairs: name=dns1.example.com.",
        required=True,
    )
    update.add_argument("-n", "--name", required=True)
    update.add_argument("-r", "--rtype", required=True)
    update.add_argument("-z", "--zone", required=True)

    argcomplete.autocomplete(parser)
    args = parser.parse_args()

    # Route based on command
    url = setup_url(baseurl, args.data, args.name, args.rtype, args.zone)
    if args.command == "add":
        run_add(url, args.json, headers)
    elif args.command == "config":
        run_config(args.baseurl, args.username, args.password)
    elif args.command == "delete":
        run_delete(url, args.json, headers)
    elif args.command == "list":
        run_list(url, args.json, headers)
    elif args.command == "update":
        run_update(url, args.json, headers)
    return 0


if __name__ == "__main__":
    sys.exit(main())