import unittest

from patchbay_web import validate

from .helpers import ED25519, AppCase


class ValidateTest(unittest.TestCase):
    def test_pubkey(self):
        self.assertEqual(validate.pubkey(ED25519), " ".join(ED25519.split()[:2]))
        for bad in ["", "ssh-ed25519", "ssh-dss AAAA", "ssh-ed25519 !!!!",
                    'ssh-rsa AAAAC3NzaC1lZDI1NTE5AAAAINJ2M1DCOTYJkGHaPO2iMGqF7isZXmwqN5HxtLP5+xq1']:
            with self.assertRaises(validate.Invalid):
                validate.pubkey(bad)

    def test_graph_rules(self):
        src = {"id": -1, "type": "client_source", "client_id": 1, "port": 22, "proto": "tcp"}
        sink = {"id": -2, "type": "public_sink", "host": "0.0.0.0", "port": 2200, "proto": "tcp"}
        split = {"id": -3, "type": "splitter"}
        nodes, links = validate.graph({"nodes": [src, sink], "links": [{"from": -1, "to": -2}]}, {1})
        self.assertEqual(links, [(-1, -2)])
        with self.assertRaises(validate.Invalid):  # wrong direction
            validate.graph({"nodes": [src, sink], "links": [{"from": -2, "to": -1}]}, {1})
        with self.assertRaises(validate.Invalid):  # unknown client
            validate.graph({"nodes": [src], "links": []}, set())
        with self.assertRaises(validate.Invalid):  # source with two outputs
            sink2 = dict(sink, id=-4)
            validate.graph({"nodes": [src, sink, sink2],
                            "links": [{"from": -1, "to": -2}, {"from": -1, "to": -4}]}, {1})
        with self.assertRaises(validate.Invalid):  # loop through splitters
            split2 = dict(split, id=-5)
            validate.graph({"nodes": [split, split2],
                            "links": [{"from": -3, "to": -5}, {"from": -5, "to": -3}]}, {1})
        with self.assertRaises(validate.Invalid):
            validate.graph({"nodes": [dict(sink, host="a b")], "links": []}, {1})

    def test_tunnel_nodes(self):
        tsrc = {"id": -1, "type": "tunnel_source", "client_id": None, "host": "10.100.0.2", "port": 25565,
                "proto": "tcp", "iface": "tun0"}
        tsink = {"id": -2, "type": "tunnel_sink", "client_id": 1, "port": 8080, "proto": "tcp", "iface": "tun1"}
        nodes, links = validate.graph({"nodes": [tsrc, tsink], "links": [{"from": -1, "to": -2}]}, {1})
        self.assertEqual(nodes[0]["iface"], "tun0")
        self.assertIsNone(nodes[0]["client_id"])
        self.assertEqual(links, [(-1, -2)])
        with self.assertRaises(validate.Invalid):
            validate.graph({"nodes": [dict(tsrc, iface="tun0; rm")], "links": []}, {1})
        # iface is ignored on non-tunnel nodes
        n, _ = validate.graph({"nodes": [{"id": -3, "type": "public_sink", "port": 1, "iface": "x"}], "links": []}, {1})
        self.assertEqual(n[0]["iface"], "")

    def test_origin(self):
        sink = {"id": -1, "type": "tunnel_sink", "port": 80, "iface": "tun0", "origin": "proxy_v2"}
        src = {"id": -2, "type": "client_source", "client_id": 1, "port": 80, "origin": "transparent"}
        nodes, _ = validate.graph({"nodes": [sink, src], "links": []}, {1})
        self.assertEqual(nodes[0]["origin"], "proxy_v2")
        self.assertEqual(nodes[1]["origin"], "")  # only sinks carry an origin
        with self.assertRaises(validate.Invalid):
            validate.graph({"nodes": [dict(sink, origin="both")], "links": []}, {1})



class GraphApiTest(AppCase):
    def test_roundtrip_keeps_ids(self):
        self.login()
        tok = self.token()
        r = self.client.post("/settings/clients/add", data={"name": "c1", "pubkey": ED25519, "csrf_token": tok})
        self.assertEqual(r.status_code, 302)
        cid = self.client.get("/api/graph").get_json()["clients"][0]["id"]

        payload = {
            "nodes": [
                {"id": -1, "type": "client_source", "client_id": cid, "host": "127.0.0.1", "port": 22, "proto": "tcp", "x": 0, "y": 0},
                {"id": -2, "type": "splitter", "x": 100, "y": 0},
                {"id": -3, "type": "public_sink", "host": "0.0.0.0", "port": 2200, "proto": "tcp",
                 "origin": "transparent", "x": 200, "y": 0},
            ],
            "links": [{"from": -1, "to": -2}, {"from": -2, "to": -3}],
        }
        r = self.client.put("/api/graph", json=payload, headers={"X-CSRF-Token": tok})
        self.assertEqual(r.status_code, 200, r.get_data(as_text=True))
        ids = r.get_json()["ids"]
        self.assertFalse(r.get_json()["daemon"])  # no daemon in tests

        g = self.client.get("/api/graph").get_json()
        self.assertEqual(len(g["nodes"]), 3)
        sink_id = ids["-3"]
        self.assertEqual(next(n for n in g["nodes"] if n["id"] == sink_id)["origin"], "transparent")
        # Saving again with real ids keeps them (stats are keyed by sink id).
        for n in payload["nodes"]:
            n["id"] = ids[str(n["id"])]
        payload["links"] = [{"from": ids["-1"], "to": ids["-2"]}, {"from": ids["-2"], "to": ids["-3"]}]
        r = self.client.put("/api/graph", json=payload, headers={"X-CSRF-Token": tok})
        self.assertEqual(r.get_json()["ids"][str(sink_id)], sink_id)

        # Deleting the client keeps its nodes but clears the reference.
        self.client.post(f"/settings/clients/{cid}/delete", data={"csrf_token": tok})
        g = self.client.get("/api/graph").get_json()
        self.assertIsNone(next(n for n in g["nodes"] if n["type"] == "client_source")["client_id"])

    def test_invalid_graph_rejected(self):
        self.login()
        tok = self.token()
        r = self.client.put("/api/graph", json={"nodes": [{"id": -1, "type": "bogus"}]}, headers={"X-CSRF-Token": tok})
        self.assertEqual(r.status_code, 400)


if __name__ == "__main__":
    unittest.main()