#!/usr/bin/python3

import sys
import re
import json
import dns.resolver
import asyncio

dkim_selectors = [ # <selector>._domainkey.<domain>
    'dkim',
    'google', # Google
    'selector', 'selector1', 'selector2', 'selector3', # Microsoft
    'fm1', 'fm2', 'fm3', 'mesmtp', # FastMail
    'protonmail', 'protonmail1', 'protonmail2', # ProtonMail
    'sg', 's1', 'sg2', 's2', 's3', # SendGrid
    'ctct1', 'ctct2', # Constant Contact
    'k1', # MailChimp
    'mandrill', # MailChimp Add-on
    'pic', 'smtp', # Mailgun (Very generic, thank you, Mailgun)
    'bh', # Bullhorn
    'cm' # Campaign Monitor
]

domains = []

def safe_resolve(name: str, type: str) -> dns.resolver.Answer:
    try:
        return dns.resolver.resolve(name, type)
    except dns.resolver.NoAnswer:
        return None
    except dns.resolver.NXDOMAIN:
        return None
    except dns.exception.Timeout:
        return None

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

    def resolve_mx(self) -> None:
        answer = safe_resolve(self.name, 'MX')
        if answer:
            for a in answer:
                self.mx.append({
                    'exchanger': re.sub(r'^\d+\s', '', str(a)),
                    'preference': re.search(r'^(\d+)', str(a)).groups()[0]
                })

    def get_mx(self) -> list:
        return self.mx
    
    # TODO: Recurse for each 'include:' in SPF
    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('"', '')
                    matches = re.search(r'(?:include:)(.*)(?=\s)', str(a).lower())
                    domains = []
                    if matches:
                        for match in matches:
                            domains.append(Domain(match.group(1)))
                        return domains
                    break
        else: self.spf = None

    def get_spf(self) -> str:
        return self.spf
    
    def resolve_dkim(self) -> None:
        for s in dkim_selectors:
            # Search for CNAME DKIM records
            answer = safe_resolve(f'{s}._domainkey.{self.name}', 'CNAME')
            if answer:
                self.dkim.append({'type': 'CNAME', 'value': str(answer[0]).replace('"', '')})
                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({'type': 'TXT', 'value': str(answer[0]).replace('"', '')})
        # Check for TXT DKIM record
        answer = safe_resolve(self.name, 'TXT')
        if answer:
            for a in answer:
                if (str(a).lower().find('v=dkim1') != -1) or (str(a).lower().find('k=rsa') != -1):
                    self.dkim.append({'type': 'TXT', 'value': str(answer[0]).replace('"', '')})

    def get_dkim(self) -> list:
        return self.dkim
    
    def resolve_dmarc(self) -> None:
        answer = safe_resolve(f'_dmarc.{self.name}', 'TXT')
        if answer: self.dmarc = str(answer[0]).replace('"', '')
        else: self.dmarc = None

    def get_dmarc(self) -> str:
        return self.dmarc
    
    def resolve_email_records(self) -> None:
        self.resolve_mx()
        self.resolve_spf()
        self.resolve_dkim() 
        self.resolve_dmarc()

if __name__ == "__main__":

    for arg in sys.argv[1:]:
        domain = Domain(arg)
        domain.resolve_email_records()
        domains.append(domain)

    print(json.dumps([domain.__dict__ for domain in domains], indent=4))