from typing import Any, Dict
import json

start_all()
machine.wait_for_unit("netbox.target")
machine.wait_until_succeeds("journalctl --since -1m --unit netbox --grep Listening")

test_objects = {
    "sites": {
        "test-site": {
            "name": "Test site",
            "slug": "test-site"
        },
        "test-site-two": {
            "name": "Test site 2",
            "slug": "test-site-second-edition"
        }
    },
    "prefixes": {
        "v4-with-updated-desc": {
            "prefix": "192.0.2.0/24",
            "class_type": "Prefix",
            "family": { "label": "IPv4" },
            "scope": {
                "__typename": "SiteType",
                "id": "1",
                "description": "Test site description"
            }
        },
        "v6-cidr-32": {
            "prefix": "2001:db8::/32",
            "class_type": "Prefix",
            "family": { "label": "IPv6" },
            "scope": {
                "__typename": "SiteType",
                "id": "1",
                "description": "Test site description"
            }
        },
        "v6-cidr-48": {
            "prefix": "2001:db8:c0fe::/48",
            "class_type": "Prefix",
            "family": { "label": "IPv6" },
            "scope": {
                "__typename": "SiteType",
                "id": "1",
                "description": "Test site description"
            }
        }
    }
}

def compare(a: str, b: str):
    differences = [(x - y) for (x,y) in list(zip(
        list(map(int, a.split('.'))),
        list(map(int, b.split('.')))
    ))]
    for d in differences:
        if d != 0:
            return d
    return 0

with subtest("Home screen loads"):
    machine.wait_until_succeeds(
        "curl -sSfL http://[::1]:80 | grep '<title>Home | NetBox</title>'"
    )

with subtest("Staticfiles are generated"):
    machine.wait_for_file("/var/lib/netbox/static/netbox.js")

with subtest("Superuser can be created"):
    machine.succeed(
        "netbox-manage createsuperuser --noinput --username netbox --email netbox@example.com"
    )
    # Django doesn't have a "clean" way of inputting the password from the command line
    machine.succeed("cat '${changePassword}' | netbox-manage shell --interface python")

def login(username: str, password: str):
    encoded_data = json.dumps({"username": username, "password": password})
    uri = "/users/tokens/provision/"
    result = json.loads(
        machine.succeed(
            "curl -sSfL "
            "-X POST "
            "-H 'Accept: application/json' "
            "-H 'Content-Type: application/json' "
            f"'http://localhost/api{uri}' "
            f"--data '{encoded_data}'"
        )
    )
    return result["key"]


netbox_version = "${netboxVersion}"
if compare(netbox_version, '4.5.2') < 0:
    with subtest("Can login"):
         auth_token = login("netbox", "netbox")

else:
    with subtest("Superusertoken can be created"):
        full_token = "0123456789abcdef0123456789abcdef01234567"
        stdout = machine.succeed("cat ${createToken} | netbox-manage shell")
        token_prefix = stdout.split(r"Bearer ")[-1].lstrip().rstrip()
        auth_token = f"{token_prefix}{full_token}"

machine.wait_for_unit("network.target")

with subtest("Home screen loads from nginx"):
    machine.succeed(
        "curl -sSfL http://localhost | grep '<title>Home | NetBox</title>'"
    )

with subtest("Staticfiles can be fetched"):
    machine.succeed("curl -sSfL http://localhost/static/netbox.js")
    machine.succeed("curl -sSfL http://localhost/static/docs/")

def get(uri: str):
    return json.loads(
        machine.succeed(
            "curl -sSfL "
            "-H 'Accept: application/json' "
            f"-H 'Authorization: Token {auth_token}' "
            f"'http://localhost/api{uri}'"
        )
    )

def delete(uri: str):
    return machine.succeed(
        "curl -sSfL "
        f"-X DELETE "
        "-H 'Accept: application/json' "
        f"-H 'Authorization: Token {auth_token}' "
        f"'http://localhost/api{uri}'"
    )


def data_request(uri: str, method: str, data: Dict[str, Any]):
    encoded_data = json.dumps(data)
    return json.loads(
        machine.succeed(
            "curl -sSfL "
            f"-X {method} "
            "-H 'Accept: application/json' "
            "-H 'Content-Type: application/json' "
            f"-H 'Authorization: Token {auth_token}' "
            f"'http://localhost/api{uri}' "
            f"--data '{encoded_data}'"
        )
    )

def post(uri: str, data: Dict[str, Any]):
    return data_request(uri, "POST", data)

def patch(uri: str, data: Dict[str, Any]):
    return data_request(uri, "PATCH", data)

with subtest("Can retrieve netbox version"):
    t.assertEqual(netbox_version, get("/status/")["netbox-version"])

with subtest("Can create objects"):
    result = post("/dcim/sites/", {"name": "Test site", "slug": "test-site"})
    site_id = result["id"]


    for prefix in test_objects["prefixes"].values():
        if compare(netbox_version, '4.2.0') >= 0:
            post("/ipam/prefixes/", {
                "prefix": prefix["prefix"],
                "scope_id": site_id,
                "scope_type": "dcim." + prefix["scope"]["__typename"].replace("Type", "").lower()
            })
            prefix["scope"]["id"] = str(site_id)
        else:
            post("/ipam/prefixes/", {
                "prefix": prefix["prefix"],
                "site": str(site_id),
            })

    result = post(
        "/dcim/manufacturers/",
        {"name": "Test manufacturer", "slug": "test-manufacturer"}
    )
    manufacturer_id = result["id"]

    # Had an issue with device-types before NetBox 3.4.0
    result = post(
        "/dcim/device-types/",
        {
            "model": "Test device type",
            "manufacturer": manufacturer_id,
            "slug": "test-device-type",
        },
    )
    device_type_id = result["id"]

with subtest("Can list objects"):
    result = get("/dcim/sites/")

    t.assertEqual(result["count"], 1)
    t.assertEqual(result["results"][0]["id"], site_id)
    t.assertEqual(result["results"][0]["name"], "Test site")
    t.assertEqual(result["results"][0]["description"], "")

    result = get("/dcim/device-types/")
    t.assertEqual(result["count"], 1)
    t.assertEqual(result["results"][0]["id"], device_type_id)
    t.assertEqual(result["results"][0]["model"], "Test device type")

with subtest("Can update objects"):
    new_description = "Test site description"
    patch(f"/dcim/sites/{site_id}/", {"description": new_description})
    result = get(f"/dcim/sites/{site_id}/")
    t.assertEqual(result["description"], new_description)

with subtest("Can delete objects"):
    # Delete a device-type since no object depends on it
    delete(f"/dcim/device-types/{device_type_id}/")

    result = get("/dcim/device-types/")
    t.assertEqual(result["count"], 0)

def request_graphql(query: str):
    return machine.succeed(
        "curl -sSfL "
        "-H 'Accept: application/json' "
        "-H 'Content-Type: application/json' "
        f"-H 'Authorization: Token {auth_token}' "
        "'http://localhost/graphql/' "
        f"--data '{json.dumps({"query": query})}'"
    )


if compare(netbox_version, '4.2.0') >= 0:
    with subtest("Can use the GraphQL API (NetBox 4.2.0+)"):
        graphql_query = '''query {
          prefix_list {
            prefix
            class_type
            family {
              label
            }
            scope {
              __typename
              ... on SiteType {
                id
                description
              }
            }
          }
        }
        '''

        answer = request_graphql(graphql_query)
        result = json.loads(answer)
        t.assertEqual(len(result["data"]["prefix_list"]), 3)
        t.assertIn(test_objects["prefixes"]["v4-with-updated-desc"], result["data"]["prefix_list"])
        t.assertIn(test_objects["prefixes"]["v6-cidr-32"], result["data"]["prefix_list"])
        t.assertIn(test_objects["prefixes"]["v6-cidr-48"], result["data"]["prefix_list"])

if compare(netbox_version, '4.2.0') < 0:
    with subtest("Can use the GraphQL API (Netbox <= 4.2.0)"):
        answer = request_graphql('''query {
          prefix_list {
            prefix
            site {
              id
            }
          }
        }
        ''')
        result = json.loads(answer)
        print(result["data"]["prefix_list"][0])
        t.assertEqual(result["data"]["prefix_list"][0]["prefix"], test_objects["prefixes"]["v4-with-updated-desc"]["prefix"])
        t.assertEqual(int(result["data"]["prefix_list"][0]["site"]["id"]), int(test_objects["prefixes"]["v4-with-updated-desc"]["scope"]["id"]))

# With 4.5.2 and higher, obtaining a session cookie or token without supplying
# proper CSRF tokens on the frontend /login/ endpoint is no longer possible
if compare(netbox_version, '4.5.2') < 0:
    with subtest("Can login with LDAP"):
        machine.wait_for_unit("openldap.service")
        login("alice", "${testPassword}")

    with subtest("Can associate LDAP groups"):
        result = get("/users/users/?username=${testUser}")

        t.assertEqual(result["count"], 1)
        t.assertTrue(any(group["name"] == "${testGroup}" for group in result["results"][0]["groups"]))

# Print systemd unit hardening state
machine.log(machine.execute("systemd-analyze security netbox.service netbox-rq.service netbox-housekeeping.service | grep -v ✓")[1])