# requires httpx, dnspython
import re, uuid, yaml, requests, argparse, urllib.parse, base64, httpx, dns.message, dns.query, dns.rdatatype



parser = argparse.ArgumentParser()
parser.add_argument("-p", "--print", action="store_true", help="print formatted device information and URLs")
parser.add_argument("--suffix", action="store", help="sets a suffix to device information, defaults to '-wan'")
parser.add_argument("-c", "--config", action="store", help="specify custom location to config.yaml")
args = parser.parse_args()

# https://stackoverflow.com/a/54254115
def isValidUUID(val):
    try:
        uuid.UUID(str(val))
        return True
    except ValueError:
        return False

# https://stackoverflow.com/a/7160778
def isValidURL(url):
    regex = re.compile(
        r'^(?:http|ftp)s?://' # http:// or https://
        r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}\.?|[A-Z0-9-]{2,}\.?)|' #domain...
        r'localhost|' #localhost...
        r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})' # ...or ip
        r'(?::\d+)?' # optional port
        r'(?:/?|[/?]\S+)$', re.IGNORECASE)
    return(re.match(regex, str(url)) is not None)

mobileConfigTemplate = """
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
  <dict>
    <key>PayloadContent</key>
    <array>
REPLACE_DNS_RECORDS
    </array>
    <key>PayloadDescription</key>
    <string>Adds encrypted DNS configurations (DNS-over-HTTPS or DNS-over-TLS) to macOS 11.0 Big Sur (or newer) and iOS 14 (or newer) based systems</string>
    <key>PayloadDisplayName</key>
    <string>Personal DoH Settings</string>
    <key>PayloadIdentifier</key>
    <string>REPLACE_IDENTIFIER</string>
    <key>PayloadRemovalDisallowed</key>
    <false/>
    <key>PayloadType</key>
    <string>Configuration</string>
    <key>PayloadUUID</key>
    <string>REPLACE_UUID</string>
    <key>PayloadVersion</key>
    <integer>1</integer>
  </dict>
</plist>
"""

EDNSRecordTemplate = """
      <dict>
        <key>DNSSettings</key>
        <dict>
          <key>DNSProtocol</key>
          <string>HTTPS</string>
          <key>ServerAddresses</key>
          <array/>
          <key>ServerURL</key>
          <string>REPLACE_SERVER_URL</string>
        </dict>
        <key>OnDemandRules</key>
        <array>
          <dict>
            <key>Action</key>
            <string>Disconnect</string>
            <key>SSIDMatch</key>
            <array>
REPLACE_EXCLUDED_SSID
            </array>
          </dict>
          <dict>
            <key>Action</key>
            <string>Connect</string>
            <key>InterfaceTypeMatch</key>
            <string>WiFi</string>
          </dict>
          <dict>
            <key>Action</key>
            <string>Connect</string>
            <key>InterfaceTypeMatch</key>
            <string>Cellular</string>
          </dict>
          <dict>
            <key>Action</key>
            <string>Disconnect</string>
          </dict>
        </array>
        <key>PayloadDescription</key>
        <string>Configures device to use REPLACE_DNS_NAME DNS-over-HTTPS</string>
        <key>PayloadDisplayName</key>
        <string>REPLACE_DNS_NAME DoH</string>
        <key>PayloadIdentifier</key>
        <string>com.apple.dnsSettings.managed.REPLACE_UUID</string>
        <key>PayloadType</key>
        <string>com.apple.dnsSettings.managed</string>
        <key>PayloadUUID</key>
        <string>REPLACE_UUID</string>
        <key>PayloadVersion</key>
        <integer>1</integer>
        <key>ProhibitDisablement</key>
        <false/>
      </dict>
"""

if (args.config):
    with open(args.config, 'r') as file:
        config = yaml.safe_load(file)
else:
    with open('config.yaml', 'r') as file:
        config = yaml.safe_load(file)

def isValidDOH(url):
    try:
        # resp = requests.get(
        #     url = url,
        #     params = {"name": "one.one.one.one", "type": "A"},
        #     headers = {"accept": "application/dns-json"}
        # )
        # for ans in resp.json()["Answer"]:
        #     if ans["data"] == "1.1.1.1" or ans["data"] == "1.0.0.1":
        #         return True

        # reqRaw = dns.message.make_query("one.one.one.one", 'A', id=0).to_wire()
        # reqStr = base64.urlsafe_b64encode(reqRaw).decode('utf-8').rstrip('=')
        # resp = requests.get(
        #     url = url,
        #     params = {"dns": reqStr},
        #     headers = {"accept": "application/dns-message"}
        # )
        # respDecoded = str(dns.message.from_wire(resp.content))

        # if ("1.1.1.1" in respDecoded or "1.0.0.1" in respDecoded):
        #     return True
        # else:
        #     print(resp)
        #     print(respDecoded)
        #     return False

        with httpx.Client(follow_redirects=True) as client:
            q = dns.message.make_query("one.one.one.one", dns.rdatatype.A)
            r = dns.query.https(q, url, session=client)
            for answer in r.answer:
                if ("1.1.1.1" in str(answer) or "1.0.0.1" in str(answer)):
                    return True
                else:
                    print(answer)
        return False
    except Exception as e:
        print(e)
        return False


def formatDeviceInfo(deviceName, deviceUUID, toFormat = True, addSuffix = True):
    if (args.suffix is not None):
        suffix = args.suffix
        suffix = re.sub(r"[^\w\s]", '', suffix)
        suffix = re.sub(r"\s+", '-', suffix)
    else:
        suffix = "wan"

    if (not isValidUUID(deviceUUID)):
        print("ERROR: Invalid device UUID!")
        print("  Device:", deviceName)
        print("  UUID:", device["uuid"])
        exit()
    if (toFormat):
        deviceName = deviceName.lower()
        # https://stackoverflow.com/a/1007615
        # Remove all non-word characters (everything except numbers and letters)
        deviceName = re.sub(r"[^\w\s]", '', deviceName)
        # Replace all runs of whitespace with a single dash
        deviceName = re.sub(r"\s+", '-', deviceName)
        if (addSuffix):
            return deviceUUID + '-' + deviceName + '-' + suffix
        else:
            return deviceUUID + '-' + deviceName
    else:
        return urllib.parse.quote(deviceName + ' (' + deviceUUID + ')', safe='')

def newUUID4(match):
    return str(uuid.uuid4())

def main():
    newRecords = []
    for device in config["devices"]:
        if (config["individualConfig"]):
            newRecords = []

        for dns in config["dns"]:
            url = dns["url"]
            if (not isValidURL(url)):
                print("ERROR: Invalid URL!")
                print("  URL:", url)
                exit()
            if (url[-1] == '/'):
                url = url[:-1]

            if (dns["includeDeviceInfo"]):
                url += '/' if (url[-1] != '/') else ''
                urlPath = formatDeviceInfo(device["name"], device["uuid"], dns["formatDeviceInfo"])
                url = url + urlPath

            if (not isValidDOH(url)):
                print("ERROR: Invalid DOH URL!")
                print("  URL:", url)
                exit()

            print(url)

            excludedSSIDs = []
            for ssid in config["excluded_ssids"]:
                excludedSSIDs.append("                <string>" + ssid + "</string>")
            excludedSSIDs = '\n'.join(excludedSSIDs)
            newRecord = EDNSRecordTemplate
            newRecord = re.sub(r"REPLACE_EXCLUDED_SSID", excludedSSIDs, newRecord)
            newRecord = re.sub(r"REPLACE_SERVER_URL", url, newRecord)

            if (dns["includeDeviceInfo"]):
                newRecord = re.sub(r"REPLACE_DNS_NAME", dns["name"] + ' (' + device["name"] + ')', newRecord)
                newRecords.append(newRecord)
            else:
                newRecord = re.sub(r"REPLACE_DNS_NAME", dns["name"], newRecord)
                if (newRecord not in newRecords):
                    newRecords.append(newRecord)

            
        if (config["individualConfig"]):
            newRecords = "\n".join(newRecords)

            mobileConfig = mobileConfigTemplate
            mobileConfig = re.sub(r"REPLACE_IDENTIFIER", config["identifier"] + '-' + formatDeviceInfo(device["name"], device["uuid"], True, False), mobileConfig)
            mobileConfig = re.sub(r"REPLACE_DNS_RECORDS", newRecords, mobileConfig)
            mobileConfig = re.sub(r"REPLACE_UUID", newUUID4, mobileConfig)
            with open(formatDeviceInfo(device["name"], device["uuid"], True, False) + '-' + "DoH.mobileconfig", "w") as file:
                file.write(mobileConfig)

    if (not config["individualConfig"]):
        newRecords = "\n".join(newRecords)

        mobileConfig = mobileConfigTemplate
        mobileConfig = re.sub(r"REPLACE_IDENTIFIER", config["identifier"], mobileConfig)
        mobileConfig = re.sub(r"REPLACE_DNS_RECORDS", newRecords, mobileConfig)
        mobileConfig = re.sub(r"REPLACE_UUID", newUUID4, mobileConfig)
        with open("DoH.mobileconfig", "w") as file:
            file.write(mobileConfig)

if (args.print):
    for device in config["devices"]:
        print(formatDeviceInfo(device["name"], device["uuid"], True))
else:
    main()