#!/usr/bin/env python

# A partial implementation of IOC "safe sharing" (IETF draft
# draft-grimminck-safe-ioc-sharing)

# Written by Google Gemini at the beginning and completed by Stéphane
# Bortzmeyer because it takes way too many iterations with the LLM to
# address every corner case.

import re
import unittest
import json

class SafeIOC:
    """
    Provides idempotent defanging and deterministic refanging of IOCs.
    """
    LC_SCHEME_MAP = { # TODO: is it really a good idea to hardcode the
                   # list of possible schemes?
        "http": None,
        "https": None,
        "ftp": None,
        "sftp": None,
        "mailto": None
    }
    SCHEME_MAP = {}
    for scheme in LC_SCHEME_MAP:
        SCHEME_MAP[scheme.upper()] = "[%s]" % scheme.upper()
        SCHEME_MAP[scheme] = "[%s]" % scheme
            
    REVERSE_SCHEME_MAP = {v: k for k, v in SCHEME_MAP.items()}

    @classmethod
    def defang(cls, ioc: str) -> str:
        if not ioc:
            return ioc

        # Percent-decode the dot
        ioc = ioc.replace("%2e", ".")
        ioc = ioc.replace("%2E", ".")
        
        # Defang known URI schemes if present
        scheme_prefix = ""
        for raw_scheme, defanged_scheme in cls.SCHEME_MAP.items():
            if ioc.startswith(f"{raw_scheme}:"):
                scheme_prefix = f"{defanged_scheme}:"
                ioc = ioc[len(raw_scheme)+1:]
                break
            elif ioc.startswith(f"{defanged_scheme}:"):
                scheme_prefix = f"{defanged_scheme}:"
                ioc = ioc[len(defanged_scheme)+1:]
                break

        # Handle Email Address syntax if no scheme prefix exists
        if not scheme_prefix and "@" in ioc and "[@]" not in ioc:
            ioc = ioc.replace("@", "[@]")

        if scheme_prefix.lower() == "[mailto]:" and "@" in ioc and "[@]" not in ioc:
            ioc = ioc.replace("@", "[@]")
            
        # http://user:passwd@host case
        if (scheme_prefix.lower() == "[http]:" or scheme_prefix.lower() == "[https]:") and "@" in ioc and "[@]" not in ioc:
            ioc = ioc.replace("@", "[@]")
            
        # Identify and isolate the host boundaries
        # This prevents accidental defanging of colons in trailing port numbers or paths
        host = ioc
        suffix = ""
        
        # Handle IPv6 inside URLs (bracketed literals)
        # Find everything between '[' and ']' occurring early in the string
        if "//[" in ioc:
            prefix, rest = ioc.split("//[", 1)
            if "]" in rest:
                ipv6_content, suffix = rest.split("]", 1)
                # Obfuscate colons inside the IPv6 bracket
                defanged_ipv6 = re.sub(r'(?<!\[)\:(?!\])', '[:]', ipv6_content)
                # Obfuscate dots inside the IPv6 bracket (for IPv4-mapped IPv6)
                defanged_ipv6 = re.sub(r'(?<!\[)\.(?!\])', '[.]', defanged_ipv6)
                
                # Reconstruct url, then let step 4 handle any non-host details if needed
                ioc = f"{defanged_scheme}://[{defanged_ipv6}]{suffix}"
                return ioc

        # If it is a raw unbracketed hybrid IPv6 address (detected by containing at least two colons and dots)
        if host.count(":") >= 2:
            # Safely split off trailing ports or paths if they exist
            # Find the end of the IPv6 block (either slash for path, or port if not bracketed)
            break_point = len(host)
            for char in ['/', '?', '#']:
                pos = host.find(char)
                if pos != -1 and pos < break_point:
                    break_point = pos
            
            suffix = host[break_point:]
            host = host[:break_point]
            
            # Apply both dot and colon transformations to the standalone host
            host = re.sub(r'(?<!\[)\:(?!\])', '[:]', host)
            host = re.sub(r'(?<!\[)\.(?!\])', '[.]', host)
            return f"{scheme_prefix}{host}{suffix}"

        # Standard Fallback for Domains/IPv4 strings
        ioc = re.sub(r'(?<!\[)\.(?!\])', '[.]', ioc)
        return f"{scheme_prefix}{ioc}"
            
    @classmethod
    def refang(cls, defanged_ioc: str) -> str:
        if not defanged_ioc:
            return defanged_ioc

        ioc = defanged_ioc
        
        # 1. Restore URI Scheme
        for defanged_scheme, raw_scheme in cls.REVERSE_SCHEME_MAP.items():
            if ioc.startswith(f"{defanged_scheme}://"):
                ioc = f"{raw_scheme}://" + ioc[len(defanged_scheme)+3:]
                break
                
        # 2. Restore Email syntax
        ioc = ioc.replace("[@]", "@")
        
        # 3. Restore standard dot syntax
        ioc = ioc.replace("[.]", ".")
        
        return ioc

class TestSafeIOCSharing(unittest.TestCase):
    
    def test_url_defanging(self):
        raw = "https://bad.example.com/path"
        expected = "[https]://bad[.]example[.]com/path"
        self.assertEqual(SafeIOC.defang(raw), expected)

    def test_idempotency(self):
        # A string that is already defanged should remain identical
        already_defanged = "[https]://bad[.]example[.]com/path"
        self.assertEqual(SafeIOC.defang(already_defanged), already_defanged)

    def test_ipv4_defanging(self):
        raw = "192.0.2.1"
        expected = "192[.]0[.]2[.]1"
        self.assertEqual(SafeIOC.defang(raw), expected)

    def test_email_defanging(self):
        raw = "phish@malicious.domain.com"
        expected = "phish[@]malicious[.]domain[.]com"
        self.assertEqual(SafeIOC.defang(raw), expected)

    def test_punycode_defanging(self):
        raw = "xn--n3h.example"
        expected = "xn--n3h[.]example"
        self.assertEqual(SafeIOC.defang(raw), expected)

    def test_full_reversibility_cycle(self):
        raw = "https://bad.example.com/path"
        defanged = SafeIOC.defang(raw)
        self.assertEqual(SafeIOC.refang(defanged), raw)

    def test_vectors_in_the_draft(self):
        data = json.loads("""
     [   
     {"label": "standard-url", "operation": "obfuscate",
      "input": "https://bad.example",
      "expected": "[https]://bad[.]example"},
     {"label": "url-with-path", "operation": "obfuscate",
      "input": "https://evil.example/path",
      "expected": "[https]://evil[.]example/path"},
     {"label": "deep-link-url", "operation": "obfuscate",
      "input": "https://bad.example/path/to/page?q=1#frag",
      "expected": "[https]://bad[.]example/path/to/page?q=1#frag"},
     {"label": "http-url", "operation": "obfuscate",
      "input": "http://attacker.example",
      "expected": "[http]://attacker[.]example"},
     {"label": "ftp-url", "operation": "obfuscate",
      "input": "ftp://files.example/",
      "expected": "[ftp]://files[.]example/"},
     {"label": "mailto", "operation": "obfuscate",
      "input": "mailto:user@example.com",
      "expected": "[mailto]:user[@]example[.]com"},
     {"label": "ipv4-address", "operation": "obfuscate",
      "input": "198.51.100.1",
      "expected": "198[.]51[.]100[.]1"},
     {"label": "ipv4-in-url", "operation": "obfuscate",
      "input": "http://192.0.2.1",
      "expected": "[http]://192[.]0[.]2[.]1"},
     {"label": "ipv6-in-url", "operation": "obfuscate",
      "input": "http://[2001:db8::1]:8080",
      "expected": "[http]://[2001[:]db8[:][:]1]:8080"},
     {"label": "ipv6-full-form", "operation": "obfuscate",
      "input": "http://[2001:db8:0:0:0:0:0:1]/",
      "expected": "[http]://[2001[:]db8[:]0[:]0[:]0[:]0[:]0[:]1]/"},
     {"label": "ipv4-mapped-ipv6", "operation": "obfuscate",
      "input": "http://[::ffff:192.0.2.1]",
      "expected": "[http]://[[:][:]ffff[:]192[.]0[.]2[.]1]"},
     {"label": "ipv4-embedded-ipv6-wkp", "operation": "obfuscate",
      "input": "http://[64:ff9b::198.51.100.1]/",
      "expected": "[http]://[64[:]ff9b[:][:]198[.]51[.]100[.]1]/"},
     {"label": "ipv4-embedded-ipv6-nsp", "operation": "obfuscate",
      "input": "2001:db8:122:344::198.51.100.1",
      "expected": "2001[:]db8[:]122[:]344[:][:]198[.]51[.]100[.]1"},
     {"label": "ipv6-with-zone-id", "operation": "obfuscate",
      "input": "http://[2001:db8::1%25eth0]/",
      "expected": "[http]://[2001[:]db8[:][:]1%25eth0]/"},
     {"label": "bare-ipv6", "operation": "obfuscate",
      "input": "2001:db8::1",
      "expected": "2001[:]db8[:][:]1"},
     {"label": "bracketed-bare-ipv6", "operation": "obfuscate",
      "input": "[2001:db8::1]",
      "expected": "[2001[:]db8[:][:]1]"},
     {"label": "bare-ipv6-with-zone", "operation": "obfuscate",
      "input": "2001:db8::1%eth0",
      "expected": "2001[:]db8[:][:]1%eth0"},
     {"label": "email-address", "operation": "obfuscate",
      "input": "phish@target.example",
      "expected": "phish[@]target[.]example"},
     {"label": "punycode-domain", "operation": "obfuscate",
      "input": "xn--n3h.example",
      "expected": "xn--n3h[.]example"},
     {"label": "url-with-userinfo", "operation": "obfuscate",
      "input": "http://user:pass@attacker.example",
      "expected": "[http]://user:pass[@]attacker[.]example"},
     {"label": "bare-domain-with-port", "operation": "obfuscate",
      "input": "evil.example:443",
      "expected": "evil[.]example:443"},
     {"label": "scheme-case-preserved", "operation": "obfuscate",
      "input": "HTTPS://bad.example",
      "expected": "[HTTPS]://bad[.]example"},
     {"label": "percent-encoded-dot-in-host", "operation": "obfuscate",
      "input": "http://evil%2eexample",
      "expected": "[http]://evil[.]example"},
     {"label": "percent-encoded-dot-upper", "operation": "obfuscate",
      "input": "http://evil%2Eexample",
      "expected": "[http]://evil[.]example"},
     {"label": "percent-encoded-at-preserved", "operation": "obfuscate",
      "input": "http://user%40attacker.example",
      "expected": "[http]://user%40attacker[.]example"},
     {"label": "percent-encoded-colon", "operation": "obfuscate",
      "input": "http://evil%3a80.example",
      "expected": "[http]://evil%3a80[.]example"},
     {"label": "percent-encoded-non-delim", "operation": "obfuscate",
      "input": "http://%65vil.example",
      "expected": "[http]://%65vil[.]example"},
     {"label": "ipv4-cidr", "operation": "obfuscate",
      "input": "192.0.2.0/24",
      "expected": "192[.]0[.]2[.]0/24"},
     {"label": "ipv6-cidr", "operation": "obfuscate",
      "input": "2001:db8::/32",
      "expected": "2001[:]db8[:][:]/32"},
     {"label": "nested-email-in-query", "operation": "obfuscate",
      "input":
        "http://example.com/?contact=abuse@evil.example",
      "expected":
        "[http]://example[.]com/?contact=abuse[@]evil[.]example"},
     {"label": "idempotency-check", "operation": "obfuscate",
      "input": "[https]://bad[.]example",
      "expected": "[https]://bad[.]example"},
     {"label": "ipv6-idempotency", "operation": "obfuscate",
      "input": "[http]://[2001[:]db8[:][:]1]:8080",
      "expected": "[http]://[2001[:]db8[:][:]1]:8080"},
     {"label": "legacy-deobfuscation", "operation": "deobfuscate",
      "input": "hxxps://bad[.]example",
      "expected": "https://bad.example"},
     {"label": "legacy-deobfuscation-bare", "operation": "deobfuscate",
      "input": "hxxps://bad.example",
      "expected": "https://bad.example"},
     {"label": "case-preserved-roundtrip", "operation": "deobfuscate",
      "input": "[HTTPS]://bad[.]example",
      "expected": "HTTPS://bad.example"},
     {"label": "ipv6-uri-deobfuscation", "operation": "deobfuscate",
      "input": "[http]://[2001[:]db8[:][:]1]:8080",
      "expected": "http://[2001:db8::1]:8080"},
     {"label": "mailto-deobfuscation", "operation": "deobfuscate",
      "input": "[mailto]:user[@]example[.]com",
      "expected": "mailto:user@example.com"}
   ]
""")
        # TODO these tests are in the draft but do not work yet
        ignored_data = json.loads("""
        [
     {"label": "nested-url-in-query", "operation": "obfuscate",
      "input":
        "http://example.com/r?url=http://evil.example",
      "expected":
        "[http]://example[.]com/r?url=[http]://evil[.]example"},
             {"label": "chained-redirect", "operation": "obfuscate",
      "input":
        "http://a.test/?u=http://b.test/?u=http://c.test",
      "expected":
        "[http]://a[.]test/?u=[http]://b[.]test/?u=[http]://c[.]test"}
        ]
""")        
        for test in data:
            if test["operation"] == "obfuscate":
                self.assertEqual(SafeIOC.defang(test["input"]), test["expected"])

if __name__ == "__main__":
    # This invokes the test runner natively
    unittest.main()