import gzip import json import os import unittest from datetime import datetime from helpers import TempDirTest from shadow_ai_report import FREE_LIST_SIZE from shadow_ai_report.parse import (FileStats, InputError, Matcher, decision_blocked, detect_format, load_tools, normalize_domain, parse_timestamp, read_log) def read_all(path, fmt="auto"): stats = FileStats(path) recs = list(read_log(path, stats, fmt)) return recs, stats class TestNormalise(unittest.TestCase): def test_case_www_trailing_dot(self): self.assertEqual(normalize_domain("WWW.Dify.AI."), "dify.ai") self.assertEqual(normalize_domain(" api.dify.ai. "), "api.dify.ai") self.assertEqual(normalize_domain("www.www.example.ai"), "www.example.ai") # one www. only self.assertEqual(normalize_domain("*.dify.ai"), "dify.ai") def test_urls_and_ports(self): self.assertEqual(normalize_domain("https://www.dify.ai/path?q=1"), "dify.ai") self.assertEqual(normalize_domain("http://API.Dify.ai:8443/x"), "api.dify.ai") self.assertEqual(normalize_domain("cloud.dify.ai:443"), "cloud.dify.ai") self.assertEqual(normalize_domain("cloud.dify.ai/some/path"), "cloud.dify.ai") self.assertEqual(normalize_domain("[2001:db8::1]:443"), "2001:db8::1") self.assertEqual(normalize_domain("192.0.2.7"), "192.0.2.7") def test_idn(self): self.assertEqual(normalize_domain("bücher.example"), "xn--bcher-kva.example") def test_invalid(self): for v in (None, "", " ", ".", "not a domain!!", "user@example.com", "a b.com", "\"\""): self.assertIsNone(normalize_domain(v), v) class TestMatching(unittest.TestCase): @classmethod def setUpClass(cls): cls.tools = load_tools() cls.m = Matcher(cls.tools) def match(self, raw): d = normalize_domain(raw) t = self.m.match(d) if d else None return t.domain if t else None def test_list_loaded(self): self.assertEqual(len(self.tools), 300) self.assertEqual(len(self.tools), FREE_LIST_SIZE) t = self.tools["dify.ai"] self.assertEqual(t.category, "Agents & Automation") self.assertEqual(t.risk_level, "low") self.assertEqual(self.tools["chadgpt.ru"].risk_level, "high") self.assertEqual(self.tools["chadgpt.ru"].data_sovereignty, "high") def test_widely_used_tools(self): self.assertEqual(self.match("chatgpt.com"), "chatgpt.com") self.assertEqual(self.match("api.openai.com"), "openai.com") self.assertEqual(self.match("https://www.perplexity.ai/search?q=x"), "perplexity.ai") self.assertEqual(self.match("CHARACTER.AI."), "character.ai") self.assertEqual(self.match("cdn.midjourney.com"), "midjourney.com") self.assertEqual(self.tools["chatgpt.com"].name, "ChatGPT") self.assertEqual(self.tools["chatgpt.com"].risk_level, "medium") def test_exact_and_subdomains(self): self.assertEqual(self.match("dify.ai"), "dify.ai") self.assertEqual(self.match("api.dify.ai"), "dify.ai") self.assertEqual(self.match("a.b.c.cloud.dify.ai"), "dify.ai") def test_www_case_trailing_dot(self): self.assertEqual(self.match("WWW.DUST.TT."), "dust.tt") self.assertEqual(self.match("Dify.Ai"), "dify.ai") self.assertEqual(self.match("https://www.poly.ai/"), "poly.ai") def test_no_false_suffix_matches(self): self.assertIsNone(self.match("notdify.ai")) # same ending, not a subdomain self.assertIsNone(self.match("dify.ai.example.com")) self.assertIsNone(self.match("ai")) self.assertIsNone(self.match("example.org")) class TestTimestamps(unittest.TestCase): def test_formats(self): want = datetime(2026, 9, 21, 8, 30, 5) for v in ("2026-09-21T08:30:05Z", "2026-09-21T08:30:05.123456789Z", "2026-09-21 08:30:05", "2026-09-21T10:30:05+02:00", "2026-09-21T03:30:05-0500", "21/Sep/2026:08:30:05 +0000", 1789979405, "1789979405", 1789979405000, 1789979405000000000, "1789979405000000000"): got = parse_timestamp(v) self.assertIsNotNone(got, v) self.assertEqual(got.replace(microsecond=0), want, v) self.assertEqual(parse_timestamp("2026-09-21"), datetime(2026, 9, 21)) def test_bad(self): for v in (None, "", "yesterday", "12", True, "2026-13-40T00:00:00Z"): self.assertIsNone(parse_timestamp(v), v) class TestDecisions(unittest.TestCase): def test_generic(self): self.assertTrue(decision_blocked("Blocked")) self.assertTrue(decision_blocked("DENY")) self.assertFalse(decision_blocked("allowed")) self.assertIsNone(decision_blocked("")) self.assertIsNone(decision_blocked("5")) def test_cloudflare(self): for v in ("blockedByCategory", "blockedAlwaysCategory", "blockedRule", 3, 6, 9): self.assertTrue(decision_blocked(v, cloudflare=True), v) for v in ("allowedOnNoPolicyMatch", "allowedOnNoLocation", "overrideForSafeSearch", "overrideApplied", "allowedRule", 5, 10): self.assertFalse(decision_blocked(v, cloudflare=True), v) class TestCsvDetection(TempDirTest): def check(self, header, row, domain="api.dify.ai", user=None, sep=","): path = self.write("log.csv", sep.join(header) + "\n" + sep.join(row) + "\n") self.assertEqual(detect_format(path), "csv") recs, stats = read_all(path) self.assertEqual(len(recs), 1, stats.as_dict()) self.assertEqual(recs[0][0], domain) self.assertEqual(recs[0][1], user) return recs[0], stats def test_header_variants(self): for name in ("domain", "Domain", "query", "query_name", "QueryName", "Query Name", "qname", "host", "hostname", "fqdn", "site", "dest", "destination", "dst_host"): self.check(["time", name], ["2026-09-21T08:00:00Z", "api.dify.ai"]) def test_user_variants(self): for name in ("user", "username", "email", "user_email", "src_user", "identity", "client", "device", "src_ip", "client_ip", "source_ip", "SrcIP", "Email"): self.check([name, "query_name"], ["someone", "api.dify.ai"], user="someone") def test_time_variants(self): for name in ("timestamp", "time", "datetime", "date", "ts", "@timestamp", "Datetime"): rec, _ = self.check([name, "domain"], ["2026-09-21T08:00:00Z", "dify.ai"], domain="dify.ai") self.assertEqual(rec[2], datetime(2026, 9, 21, 8)) def test_url_column(self): rec, stats = self.check(["timestamp", "user", "url"], ["2026-09-21T08:00:00Z", "ann", "https://www.dify.ai/apps?x=1,2"], domain="dify.ai", user="ann") self.assertEqual(stats.columns["domain"], "url") def test_host_preferred_over_url(self): rec, stats = self.check(["url", "host"], ["https://example.org/", "cloud.dify.ai"], domain="cloud.dify.ai") self.assertEqual(stats.columns["domain"], "host") def test_tsv_and_semicolon(self): _, stats = self.check(["ts", "client_ip", "qname"], ["1789979405", "192.0.2.5", "dify.ai"], domain="dify.ai", user="192.0.2.5", sep="\t") self.assertIn("tab", stats.format) self.check(["date", "user", "domain"], ["2026-09-21", "bob", "dust.tt"], domain="dust.tt", user="bob", sep=";") def test_user_fallback_to_ip(self): path = self.write("log.csv", "user,src_ip,domain\n,192.0.2.9,dify.ai\nann,192.0.2.10,dify.ai\n") recs, stats = read_all(path) self.assertEqual([r[1] for r in recs], ["192.0.2.9", "ann"]) self.assertEqual(stats.columns["user"], "user / src_ip") def test_quoted_fields_and_bom(self): path = self.write("log.csv", '"time","user","domain"\n"2026-09-21 08:00:00","Doe, Jane","dify.ai"\n') recs, _ = read_all(path) self.assertEqual(recs[0][:2], ("dify.ai", "Doe, Jane")) def test_text_list(self): path = self.write("domains.txt", "# exported list\nDify.ai.\n\nwww.dust.tt # comment\nexample.org\n" "not a domain!!\n") self.assertEqual(detect_format(path), "text") recs, stats = read_all(path) self.assertEqual([r[0] for r in recs], ["dify.ai", "dust.tt", "example.org"]) self.assertEqual((stats.lines, stats.records, stats.skipped), (4, 3, 1)) def test_headerless_csv_guessed(self): path = self.write("noheader.csv", "".join( "2026-09-21T08:%02d:00Z,user%d@example.com,api.dify.ai,A\n" % (i, i) for i in range(10))) self.assertEqual(detect_format(path), "csv") recs, stats = read_all(path) self.assertEqual(len(recs), 10) self.assertEqual(recs[0][0], "api.dify.ai") self.assertEqual(recs[0][1], "user0@example.com") self.assertIsNotNone(recs[0][2]) self.assertTrue(any("guessed" in n for n in stats.notes)) def test_no_domain_column(self): path = self.write("x.csv", "a,b,c\n1,2,3\n4,5,6\n") with self.assertRaises(InputError): read_all(path) def test_utf16(self): path = self.write("u16.csv", "domain,user\ndify.ai,ann\n", encoding="utf-16-le") recs, _ = read_all(path) self.assertEqual(recs[0][:2], ("dify.ai", "ann")) class TestMalformed(TempDirTest): def test_skipped_and_counted(self): text = ("timestamp,user,domain\n" "2026-09-21T08:00:00Z,ann,dify.ai\n" "2026-09-21T08:00:01Z,ann\n" # too few fields "2026-09-21T08:00:02Z,ann,not a domain!!\n" # invalid domain "2026-09-21T08:00:03Z,ann,\n" # empty domain '2026-09-21T08:00:04Z,"ann,dify.ai\n' # unterminated quote "\n" # blank: ignored "2026-09-21T08:00:05Z,bob,dust.tt\n" "bad-time,bob,dust.tt\n" # counted, time unreadable "2026-09-21T08:00:06Z,bob\x00,poly.ai\n") # NUL byte removed path = self.write("bad.csv", text) recs, stats = read_all(path) self.assertEqual([r[0] for r in recs], ["dify.ai", "dust.tt", "dust.tt", "poly.ai"]) self.assertEqual(stats.lines, 8) self.assertEqual(stats.records, 4) self.assertEqual(stats.skipped, 4) self.assertEqual(stats.lines, stats.records + stats.skipped) self.assertEqual(stats.bad_time, 1) def test_stray_quote_does_not_swallow_file(self): rows = ['2026-09-21T08:00:00Z,"ann,dify.ai'] + ["2026-09-21T08:00:%02dZ,bob,dust.tt" % i for i in range(20)] path = self.write("q.csv", "time,user,domain\n" + "\n".join(rows) + "\n") recs, stats = read_all(path) self.assertEqual(len(recs), 20) self.assertEqual(stats.skipped, 1) class TestCloudflare(TempDirTest): def lines(self, *objs): out = [] for o in objs: out.append(o if isinstance(o, str) else json.dumps(o)) return "\n".join(out) + "\n" def test_fields(self): text = self.lines( {"Datetime": "2026-09-21T08:00:00Z", "QueryName": "api.dify.ai", "Email": "ann@example.com", "SrcIP": "192.0.2.1", "ResolverDecision": "allowedOnNoPolicyMatch"}, {"Datetime": 1789979405000000000, "QueryName": "CHADGPT.RU.", "Email": "non_identity@team.cloudflareaccess.com", "SrcIP": "192.0.2.200", "ResolverDecision": "blockedByCategory"}, {"Datetime": "2026-09-21T08:30:05Z", "QueryName": "example.org", "Email": "", "SrcIP": "192.0.2.3", "ResolverDecision": "overrideApplied"}, '{"Datetime":"2026-09-21T09:00:00Z","QueryName":"dify.ai","Email":"bo', # truncated {"Datetime": "2026-09-21T09:00:00Z", "QueryName": ".", "Email": "ann@example.com"}, # root query {"Datetime": "2026-09-21T09:00:00Z", "Email": "ann@example.com"}, # no QueryName "[1, 2]", ) path = self.write("gw.jsonl", text) self.assertEqual(detect_format(path), "cloudflare-dns") recs, stats = read_all(path) self.assertEqual([r[0] for r in recs], ["api.dify.ai", "chadgpt.ru", "example.org"]) self.assertEqual([r[1] for r in recs], ["ann@example.com", "192.0.2.200", "192.0.2.3"]) self.assertEqual(recs[0][2], datetime(2026, 9, 21, 8)) self.assertEqual(recs[1][2], datetime(2026, 9, 21, 8, 30, 5)) self.assertEqual([r[3] for r in recs], [False, True, False]) self.assertEqual((stats.lines, stats.records, stats.skipped), (7, 3, 4)) self.assertEqual(stats.columns, {"domain": "QueryName", "user": "Email / SrcIP", "time": "Datetime", "action": "ResolverDecision"}) def test_minimal_fields(self): path = self.write("gw.json", self.lines({"QueryName": "dust.tt"}, {"QueryName": "example.org"})) recs, stats = read_all(path) self.assertEqual(recs, [("dust.tt", None, None, None), ("example.org", None, None, None)]) self.assertEqual(stats.columns, {"domain": "QueryName"}) def test_forced_format_on_wrong_file(self): path = self.write("log.csv", "domain\ndify.ai\n") recs, stats = read_all(path, fmt="cloudflare-dns") self.assertEqual(recs, []) self.assertEqual(stats.skipped, 2) def test_other_json_rejected(self): path = self.write("other.jsonl", self.lines({"host": "dify.ai"})) with self.assertRaises(InputError): detect_format(path) class TestGzip(TempDirTest): def test_gzip_csv_and_jsonl(self): p1 = self.path("log.csv.gz") with gzip.open(p1, "wt", encoding="utf-8", newline="") as f: f.write("timestamp,user,query_name\n2026-09-21T08:00:00Z,ann,cloud.dify.ai\n2026-09-21T08:00:01Z,bob,x.org\n") recs, stats = read_all(p1) self.assertEqual([r[0] for r in recs], ["cloud.dify.ai", "x.org"]) self.assertEqual(stats.format, "csv") # gzip is detected from the content, not the file name p2 = self.path("gateway_dns_20260921.log") with gzip.open(p2, "wt", encoding="utf-8") as f: f.write(json.dumps({"QueryName": "dust.tt", "Email": "ann@example.com"}) + "\n") recs, stats = read_all(p2) self.assertEqual(stats.format, "cloudflare-dns") self.assertEqual(recs[0][:2], ("dust.tt", "ann@example.com")) def test_truncated_gzip(self): raw = gzip.compress(("domain\n" + "dify.ai\n" * 20000).encode()) path = self.path("cut.csv.gz") with open(path, "wb") as f: f.write(raw[: len(raw) // 2]) recs, stats = read_all(path) self.assertGreater(len(recs), 0) self.assertTrue(any("stopped" in n for n in stats.notes), stats.notes) def test_missing_and_empty(self): with self.assertRaises(InputError): read_all(self.path("nope.csv")) with self.assertRaises(InputError): read_all(self.write("empty.csv", "")) with self.assertRaises(InputError): read_all(self.tmp) if __name__ == "__main__": unittest.main()