#!/usr/bin/python3

import sys
from dns import resolver
import re
import json

dkim_selectors = { # <selector>._domainkey.<domain>
    # Applications with dynamic selector names will not be detected, e.g. smtp2go (s12345._domainkey.<domain>)
    'not specified': [
        'dkim', 'default'
    ],
    'google': [
        'google'
    ],
    'exchange': [
        'selector','selector1', 'selector2', 'selector3'
    ],
    'fastmail': [
        'fm1', 'fm2', 'fm3', 'mesmtp'
    ],
    'protonmail': [
        'protonmail', 'protonmail1', 'protonmail2'
    ],
    'sendgrid': [
        'sg', 's1', 'sg2', 's2', 's3'
    ],
    'constant-contact': [
        'ctct1', 'ctct2'
    ],
    # MailChimp - Usually CNAME to dkimX.mcsv.net
    'mailchimp': [
        'k1', 'k2', 'k3'
    ],
    # MailChimp Add-on
    'mandrill': [
        'mandrill'
    ],
    'mailgun': [
        'pic', 'smtp', # Mailgun (Very generic, thank you, Mailgun)
    ],
    'bullhorn': [
        'bh'
    ],
    'campaign-monitor': [
     'cm'
    ],
    'presswise': [
        'presswise'
    ]
}

def process_args(args: list) -> list:
    domains = []
    args.pop(0)

    for arg in args:
        domains.append(Domain(name=arg))
    
    return domains

def safe_resolve(name: str, type: str) -> resolver.Answer:
    try:
        answer = resolver.resolve(name, type)
        return answer
    # We don't necessarily care if the name doesn't exist
    except resolver.NXDOMAIN:
        return None
    except resolver.NoAnswer:
        return None

class Domain:
    # def __init__(self, name: str, mx: str = [], spf: str = None, dkim: list = [], dmarc: str = None) -> None:
    def __init__(self, name: str) -> None:
        self.name = name
        self.mx = []
        self.spf = None
        self.dkim = []
        self.dmarc = None

    def resolve_mx(self) -> None:
        # Attempt to resolve the MX record for the domain
        answer = safe_resolve(self.name, 'MX')
        if answer:
            for a in answer:
                # Loop through each MX record, record its data and priority
                self.mx.append({
                    'exchanger': re.sub(r'^\d+\s', '', str(a)),
                    'priority': re.search(r'^(\d+)', str(a)).groups()[0]
                })

    def resolve_spf(self) -> None:
        answer = safe_resolve(self.name, 'TXT')
        if answer:
            for a in answer:
                if str(a).lower().find('v=spf1') != -1:
                    self.spf = str(a).replace('"', '')
                    # TODO: Recurse through includes
                    # matches = re.search(r'(?:include:)(.*)(?=\s)', str(a).lower())
                    # domains = set()
                    # if matches:
                    #     for match in matches:
                    #         domains.add(Domain(match.group(1)))
                    #     return domains
                    break
        else: self.spf = None
    
    def resolve_dkim(self) -> None:
        # Loop through our pre-defined selectors list, an attempt to resolve them all
        for app, selectors in dkim_selectors.items():
            for s in selectors:
                # Search for CNAME DKIM records
                answer = safe_resolve(f'{s}._domainkey.{self.name}', 'CNAME')
                if answer:
                    self.dkim.append({'application': app, 'selector': s, 'type': 'CNAME', 'value': str(answer[0]).replace('"', '')})
                    # This is a CNAME, we don't have to check if it's a TXT
                    continue
                # Search for TXT DKIM records
                answer = safe_resolve(f'{s}._domainkey.{self.name}', 'TXT')
                if answer:
                    if (str(answer[0]).lower().find('v=dkim1') != -1) or (str(answer[0]).lower().find('k=rsa') != -1):
                        self.dkim.append({'application': app, 'selector': s, 'type': 'TXT', 'value': str(answer[0]).replace('"', '')})
            # Check for TXT DKIM record at root
            answer = safe_resolve(self.name, 'TXT')
            if answer:
                for a in answer:
                    # Loop through all root TXT records
                    if (str(a).lower().find('v=dkim1') != -1) or (str(a).lower().find('k=rsa') != -1):
                        # If the record data includes `v=dkim1` then record it
                        self.dkim.append({'application': app, 'selector': None, 'type': 'TXT', 'value': str(answer[0]).replace('"', '')})
    
    def resolve_dmarc(self) -> None:
        # Attempt to resolve the DMARC record for the domain
        answer = safe_resolve(f'_dmarc.{self.name}', 'TXT')
        if answer: self.dmarc = str(answer[0]).replace('"', '')
        else: self.dmarc = None
    
    def resolve(self) -> None:
        # Resolve all records
        self.resolve_mx()
        self.resolve_spf()
        self.resolve_dkim()
        self.resolve_dmarc()



def main():
    domains = process_args(sys.argv)

    for domain in domains:
        domain.resolve()

    # Pretty print Domain dict as JSON
    print(json.dumps([domain.__dict__ for domain in domains], indent=4))

if __name__ == "__main__":
    
    main()