WWWWWW

Current Directory: /opt/dedrads
Viewing File: /opt/dedrads/update_spf
#!/usr/lib/rads/venv/bin/python3
"""SPF record updater"""
# Vanessa Vasile, InMotion Hosting, Oct 2017
# ref: https://trac.imhtech.net/T3/ticket/18096

from argparse import ArgumentParser, Action, ArgumentTypeError as BadArg
import platform
import sys
from typing import Literal
from subprocess import CalledProcessError, check_output
from collections import OrderedDict
from cpapis import whmapi1, cpapi2, CpAPIError
from rads import cpuser_safe
from rads.color import red, yellow, green, magenta, blue

HOSTNAME = platform.node()
SPAMEXP_HOSTNAME = 'smtp.servconfig.com'


def get_args() -> (
    tuple[Literal['create', 'update', 'delete'], bool, str, list[str]]
):
    """Arg Parser"""
    parser = ArgumentParser(description='Configures SPF records')
    actions_group = parser.add_mutually_exclusive_group(required=True)
    targets_group = parser.add_mutually_exclusive_group(required=True)
    # fmt: off
    actions_group.add_argument(
        '--create', '-c', action='store_const', const='create', dest='action',
        help='Creates a new SPF record and applies it. '
        'Will overwrite an existing one.',
    )
    actions_group.add_argument(
        '--update', '-r', action='store_const', const='update', dest='action',
        help='Updates the existing SPF record to be correct',
    )
    parser.add_argument(
        '--show-only', '-s', action='store_true',
        help='Do not apply the new record, just show it',
    )
    actions_group.add_argument(
        '--delete', '-x', action='store_const', const='delete', dest='action',
        help='Deletes an SPF TXT record',
    )
    targets_group.add_argument(
        '--user', '-u', action=SetUserAction,
        help='Update all domains (main, addon, parked) for this user',
    )
    targets_group.add_argument(
        '--domain', '-d', action=SetDomainAction, dest='domains',
        help='Domain to update',
    )
    # fmt: on
    args = parser.parse_args()
    return args.action, args.show_only, args.user, args.domains


class SetUserAction(Action):
    def __call__(self, parser, namespace, values, option_string=None):
        user: str = values
        if not cpuser_safe(user):
            raise BadArg(f"{user} is not a valid cPanel user or is restricted")
        try:
            api = cpapi2('DomainLookup::getbasedomains', user)
            domains = [x['domain'] for x in api['cpanelresult']['data']]
        except Exception as exc:
            print(type(exc).__name__, exc, file=sys.stderr)
            raise BadArg("Error getting domain list from cpapi2") from exc
        setattr(namespace, 'domains', domains)
        setattr(namespace, 'user', user)


class SetDomainAction(Action):
    def __call__(self, parser, namespace, values, option_string=None):
        domain: str = values
        user = whoowns(domain)
        if not user:
            raise BadArg(f"{domain} is not owned by a user on the system")
        if not cpuser_safe(user):
            raise BadArg(f"{user} is a restricted user")
        setattr(namespace, 'domains', [domain])


def err_msg(msg: str) -> None:
    print(red(msg), file=sys.stderr)


def main():
    action, show_only, user, domains = get_args()
    spf_data: dict[str, tuple[int, str]] = {}
    for domain in domains:
        spf_data.update(get_spf(domain))
    if user is None and not spf_data and action != 'create':
        # if user is None, --domain / -d was used
        err_msg(
            f"{domains[0]} does not have an SPF record. "
            "Use --create/-c instead"
        )
        sys.exit(1)
    mail_ips, domain_mail_ips = get_mail_ips(list(spf_data.keys()))
    if action == 'update':
        for domain, (line_num, txtdata) in spf_data.items():
            new_record = generate_spf(
                existing_spf=txtdata,
                mail_ips=mail_ips,
                domain=domain,
                domain_ip=domain_mail_ips.get(domain, None),
                overwrite=False,
            )
            print(blue(f"{domain}:"))
            print(magenta(f'\tExisting: "{txtdata}"'))
            if show_only:
                print(yellow(f'\tNew: "{new_record}"'))
                continue
            print(green(f'\tNew: "{new_record}"'))
            update_zone(domain, domain, line_num, new_record)
    elif action == 'create':
        for domain in domains:
            line_num, txtdata = spf_data.get(domain, (None, None))
            new_record = generate_spf(
                existing_spf=None,
                mail_ips=mail_ips,
                domain=domain,
                domain_ip=domain_mail_ips.get(domain, None),
                overwrite=True,
            )
            print(blue(f"{domain}:"))
            print(magenta(f'\tExisting: "{txtdata}"'))
            if show_only:
                print(yellow(f'\tNew: "{new_record}"'))
                continue
            print(green(f'\tNew: "{new_record}"'))
            if txtdata is None:  # add new record
                update_zone(domain, domain, None, new_record)
            else:  # update existing
                update_zone(domain, domain, line_num, new_record)
    elif action == 'delete':
        for domain, (line_num, txtdata) in spf_data.items():
            print(blue(f"{domain}:"))
            print(magenta(f'\tExisting: "{txtdata}"'))
            if show_only:
                continue
            delete_spf(domain, txtdata)
    else:
        raise RuntimeError(f"{action=}")


def get_mail_ips(domains: list[str]) -> tuple[list[str], dict[str, str]]:
    """Pulls relevant data"""
    mail_ips = []
    for path in ('/var/cpanel/mainip', '/var/cpanel/mainips/root'):
        try:
            with open(path, encoding='utf-8') as file:
                for line in file:
                    mail_ips.append(line.strip())
        except OSError:
            pass
    domain_mail_ips = {}
    try:
        with open('/etc/mailips', encoding='utf-8') as file:
            for line in file:
                dom, ipaddr = line.strip().split(':')
                if dom in domains:
                    domain_mail_ips[dom] = ipaddr.strip()
    except ValueError:
        err_msg('/etc/mailips is incorrectly formatted')
        sys.exit(1)
    except OSError:
        pass
    return mail_ips, domain_mail_ips


def whoowns(domain: str) -> str:
    """Obtain the cPanel user owning a domain"""
    try:
        user = check_output(
            ['/usr/local/cpanel/scripts/whoowns', domain], encoding='utf-8'
        )
    except CalledProcessError:
        return ''
    return user.strip()


def get_spf(domain: str) -> dict[str, tuple[int, str]]:
    """Gets the SPF record for a domain"""
    try:
        api = whmapi1('dumpzone', {"domain": domain}, check=True)
    except CpAPIError as exc:
        err_msg(str(exc))
        return {}
    spf_records = {}
    for data in api['data']['zone'][0]['record']:
        data: dict
        record_type: str = data.get('type', '')
        txtdata: str = data.get('txtdata', '')
        if 'TXT' not in record_type or 'v=spf1' not in txtdata:
            continue
        line: int = data['Line']
        name: str = data['name'].rstrip('.')

        # If we have a duplicate key, error out. This is invalid.
        if name in spf_records:
            err_msg(
                f"Duplicate SPF record for {name} Resolve this manually, "
                "then re-run the utility"
            )
            sys.exit(1)
        spf_records[name] = (line, txtdata)
    return spf_records


def update_zone(domain: str, name: str, line: int, content: str) -> None:
    """Updates the DNS zone"""
    if not name.endswith('.'):
        name = name + '.'
    args = {
        'domain': domain,
        'name': name,
        'class': 'IN',
        'ttl': 900,
        'type': 'TXT',
        'txtdata': content,
    }
    try:
        if line is not None:  # We're updating an existing record
            whmapi1('editzonerecord', args | {'line': line}, check=True)
        else:  # We're creating a new zone record
            whmapi1('addzonerecord', args, check=True)
    except CpAPIError as exc:
        print(red(f'\tError updating zone: {exc}'))


def delete_spf(domain: str, line: str) -> None:
    """Deletes an SPF record"""
    try:
        whmapi1('removezonerecord', {'zone': domain, 'line': line}, check=True)
    except CpAPIError as exc:
        print(red(f'\tError updating zone: {exc}'))
    else:
        print(green('\tSPF TXT record removed'))


def generate_spf(
    *,
    mail_ips: list[str],
    existing_spf: str,
    domain: str,
    domain_ip: str | None,
    overwrite: bool,
):
    """Generates the SPF record"""
    # The only difference between 'creating' and 'updating' a record is
    # 'updating' will ignore stuff that's already in the record,
    # unless it's wrong

    if overwrite or existing_spf is None:
        # If we're overwriting the record, we don't really care what's in it
        # This is the default record to include SE and server hostname
        # (all mail ips on this server)
        if domain_ip:
            new_record = 'v=spf1 +a +mx +a:{} +a:{} +ip4:{} ~all'.format(
                HOSTNAME,
                SPAMEXP_HOSTNAME,
                domain_ip,
            )
        else:
            new_record = 'v=spf1 +a +mx +a:{} +a:{} ~all'.format(
                HOSTNAME,
                SPAMEXP_HOSTNAME,
            )
        return new_record

    # If we have a record already, split it up so it's easier to parse
    parts = existing_spf.split()

    # Now we have to further split each part into sections
    # See: http://www.openspf.org/SPF_Record_Syntax

    mechanisms = ['+', '-', '~', '?']
    data = OrderedDict()
    index = 0
    for part in parts:

        if part.startswith('v='):
            data['version'] = {'mech': '', 'qual': '', 'val': part}
            index += 1
            continue

        # Check for modifiers
        if part.startswith('redirect='):
            _, val = part.split('=')
            if val != '':
                # If we have a modifier, we can basically just generate a
                # new record
                new_record = 'v=spf1 redirect=%s' % (val)
                return new_record
            # If we're missing a value though, remove this entirely
            continue

        if part.startswith('exp='):
            _, val = part.split('=')
            if val != '':
                data['exp'] = {'mech': '', 'qual': '', 'val': 'exp=%s' % val}
            continue

        # If the part starts with a mechanism...
        if any(part.startswith(m) for m in mechanisms):
            mech = part[0]
            part = part.lstrip(part[0])
        else:
            mech = ''

        # If the part contains a ":", add a value. Otherwise replace with a ''
        if ':' in part:
            qual, val = part.split(':')
        else:
            qual = part
            val = ''

        data[index] = {'mech': mech, 'qual': qual, 'val': val}
        index += 1

    # This is where we actually generate the record
    present = {
        'version': False,
        'server_host': False,
        'domain_mx': False,
        'domain_a': False,
        'domain_mailip': False,
        'se_host': False,
        'pass_fail': False,
    }
    delete = []

    for key, value in data.items():

        # ipv4 that matches mailips or domain ip
        if value['qual'] == 'ip4':
            if value['val'] in mail_ips:
                delete.append(key)
            if domain_ip is None:
                # we don't need this value if the domain uses a shared IP
                present['domain_mailip'] = True
            else:
                if value['val'] == domain_ip:
                    present['domain_mailip'] = True

        # Other servers' hostnames of ours
        if value['val'] not in (HOSTNAME, SPAMEXP_HOSTNAME):
            if (
                'inmotionhosting.com' in value['val']
                or 'webhostinghub.com' in value['val']
                or 'servconfig.com' in value['val']
            ):
                delete.append(key)

        # version
        if value['qual'] == '' and "v=" in value['val']:
            present['version'] = True

        # domain, server, & SE hostnames
        if value['qual'] == 'a':
            if value['val'] == HOSTNAME:
                present['server_host'] = True
            if value['val'] == SPAMEXP_HOSTNAME:
                present['se_host'] = True
            if value['val'] == domain or value['val'] == '':
                present['domain_a'] = True

        # The domain's MX
        if value['qual'] == 'mx':
            if value['val'] == domain or value['val'] == '':
                present['domain_mx'] = True

        # action
        if value['qual'] == "all":
            present['pass_fail'] = True
            data['pass'] = value
            if isinstance(key, int):  # delete the numerical key
                delete.append(key)

    for item in delete:
        data.pop(item)

    # Add anything that's missing
    if not present['domain_mailip'] and domain_ip is not None:
        data[index] = {'mech': '+', 'qual': 'a', 'val': domain_ip}
        index += 1
    if not present['version']:
        data['version'] = {'mech': '', 'qual': '', 'val': 'v=spf1'}
        index += 1
    if not present['server_host']:
        data[index] = {'mech': '+', 'qual': 'a', 'val': HOSTNAME}
        index += 1
    if not present['se_host']:
        data[index] = {'mech': '+', 'qual': 'a', 'val': SPAMEXP_HOSTNAME}
        index += 1
    if not present['domain_a']:
        data[index] = {'mech': '+', 'qual': 'a', 'val': ''}
        index += 1
    if not present['domain_mx']:
        data[index] = {'mech': '+', 'qual': 'mx', 'val': ''}
        index += 1
    if not present['pass_fail']:
        data['pass'] = {'mech': '~', 'qual': 'all', 'val': ''}
        index += 1

    # Now put all the parts together to make a new SPF record

    new_map = []
    new_map.append(data['version'])  # version first
    for key, val in sorted(data.items()):
        if isinstance(key, int):
            new_map.append(val)

    # The pass/fail and exp modifier should be last
    try:
        new_map.append(data['pass'])
    except KeyError:
        pass

    try:
        new_map.append(data['exp'])
    except KeyError:
        pass

    # Now put all the SPF elements together
    data = []  # reuse this one
    for item in new_map:
        if item['mech'] != '':
            if item['qual'] != '':
                line = '{}{}'.format(item['mech'], item['qual'])
            else:
                line = item['mech']
        else:
            line = '%s' % item['qual']

        if item['val'] != '' and item['qual'] != '':
            line = '{}:{}'.format(line, item['val'])
        else:
            line = '{}{}'.format(line, item['val'])

        data.append(line.strip())

    new_record = ' '.join(data)

    return new_record


if __name__ == "__main__":
    main()
F1l3 Manager
