#!/usr/bin/env python3
"""Import sender lists exported from Cisco ESA (IronPort) into Spamorez through the panel API.

Reads plain text files, one entry per line (commas and spaces also separate entries, '#' starts a comment):
  --blacklist FILE   HAT BLACKLIST entries   -> black list
  --whitelist FILE   HAT WHITELIST entries   -> white list
  --domains   FILE   domains from the RAT    -> domains (needs --smtpserver)

ESA entry            ->  Spamorez entry
  10.1.1.5            ->  10.1.1.5            (ip)
  10.1.1.0/24         ->  10.1.1.0/24         (ip)
  example.com         ->  @example.com        (domain)
  .example.com        ->  @example.com and *.example.com
  user@example.com    ->  user@example.com    (personal)
SBRS and DNS-list rules cannot be imported and are skipped.

Nothing is sent without --apply (dry run by default). Try --limit 1 first and check the panel.
Only the Python 3 standard library is used.

Environment: SPZ_API_URL (default https://my.spamorez.ru/panelv32/api/), SPZ_TOKEN (access token from the panel).
Example:
  export SPZ_TOKEN=...
  python3 esa_import.py --blacklist blacklist.txt --whitelist whitelist.txt            # dry run
  python3 esa_import.py --blacklist blacklist.txt --whitelist whitelist.txt --apply
"""
import argparse
import ipaddress
import json
import os
import re
import sys
import time
import urllib.parse
import urllib.request

PUBLIC_DOMAINS = {
    'gmail.com', 'googlemail.com', 'yandex.ru', 'yandex.com', 'ya.ru', 'mail.ru', 'bk.ru', 'list.ru', 'inbox.ru',
    'rambler.ru', 'outlook.com', 'hotmail.com', 'live.com', 'yahoo.com', 'icloud.com', 'me.com', 'aol.com',
}
EMAIL = re.compile(r'^[^@\s]+@[^@\s]+\.[^@\s]+$')
DOMAIN = re.compile(r'^[a-z0-9]([a-z0-9-]*[a-z0-9])?(\.[a-z0-9]([a-z0-9-]*[a-z0-9])?)+$', re.I)


def read_entries(path):
    """Return the raw entries of a text file."""
    out = []
    with open(path, encoding='utf-8', errors='replace') as f:
        for line in f:
            line = line.split('#', 1)[0].strip()
            if not line:
                continue
            if re.search(r'\b(sbrs|dnslist)\b', line, re.I):
                print('skip (SBRS/DNS list rule is not importable): ' + line, file=sys.stderr)
                continue
            out += [x for x in re.split(r'[,\s;]+', line) if x]
    return out


def convert(entry):
    """ESA entry -> list of (address, kind); kind is ip, domain or personal. Unknown entries give an empty list."""
    e = entry.strip().strip('"\'').lower()
    try:
        net = ipaddress.ip_network(e, strict=False)
        return [(str(net) if '/' in e else e, 'ip')]
    except ValueError:
        pass
    if EMAIL.match(e):
        return [(e, 'personal')]
    if e.startswith('.') and DOMAIN.match(e[1:]):
        return [('@' + e[1:], 'domain'), ('*.' + e[1:], 'domain')]
    if e.startswith('@') and DOMAIN.match(e[1:]):
        return [(e, 'domain')]
    if e.startswith('*.') and DOMAIN.match(e[2:]):
        return [(e, 'domain')]
    if DOMAIN.match(e):
        return [('@' + e, 'domain')]
    return []


def prepare(path, whitelist):
    seen, rows, bad = set(), [], []
    for raw in read_entries(path):
        conv = convert(raw)
        if not conv:
            bad.append(raw)
            continue
        for addr, kind in conv:
            dom = addr.split('@')[-1].lstrip('*.')
            if whitelist and dom in PUBLIC_DOMAINS:
                print('skip (public domain must not be whitelisted): ' + addr, file=sys.stderr)
                continue
            if addr not in seen:
                seen.add(addr)
                rows.append((addr, kind))
    return rows, bad


def api(url, token, payload):
    req = urllib.request.Request(
        url + ('&' if '?' in url else '?') + urllib.parse.urlencode({'token': token}),
        data=json.dumps(payload).encode(), headers={'Content-Type': 'application/json'})
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.loads(r.read().decode())


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('--blacklist')
    ap.add_argument('--whitelist')
    ap.add_argument('--domains')
    ap.add_argument('--smtpserver', help='target mail server for the imported domains (required with --domains)')
    ap.add_argument('--wl-status', default='any', help='white list rule type: any, domain or personal (default any)')
    ap.add_argument('--wl-to', default='(для всех)', help='recipient for white list records (default "(для всех)")')
    ap.add_argument('--comment', default='из ESA')
    ap.add_argument('--limit', type=int, default=0, help='send only the first N records of every list')
    ap.add_argument('--delay', type=float, default=0.1, help='pause between requests, seconds')
    ap.add_argument('--apply', action='store_true', help='really send the records (default: dry run)')
    a = ap.parse_args()

    url = os.environ.get('SPZ_API_URL', 'https://my.spamorez.ru/panelv32/api/')
    token = os.environ.get('SPZ_TOKEN', '')
    if a.apply and not token:
        sys.exit('SPZ_TOKEN is not set')
    if a.domains and a.apply and not a.smtpserver:
        sys.exit('--smtpserver is required with --domains')
    if not (a.blacklist or a.whitelist or a.domains):
        ap.error('give at least one of --blacklist, --whitelist, --domains')

    jobs = []   # (title, payload)
    if a.blacklist:
        rows, bad = prepare(a.blacklist, False)
        for addr, kind in rows:
            jobs.append(('blacklist ' + addr, {'request': 'blacklists', 'action': 'add', 'address': addr,
                                               'to_address': '(для всех)', 'status': kind, 'comment': a.comment}))
        for b in bad:
            print('unrecognized in blacklist: ' + b, file=sys.stderr)
    if a.whitelist:
        rows, bad = prepare(a.whitelist, True)
        for addr, kind in rows:
            jobs.append(('whitelist ' + addr, {'request': 'whitelists', 'action': 'add', 'address': addr,
                                               'to_address': a.wl_to, 'status': a.wl_status, 'comment': a.comment}))
        for b in bad:
            print('unrecognized in whitelist: ' + b, file=sys.stderr)
    if a.domains:
        for raw in read_entries(a.domains):
            d = raw.lower().lstrip('@')
            if DOMAIN.match(d):
                jobs.append(('domain ' + d, {'request': 'domain', 'action': 'add', 'domain': d,
                                             'smtpserver': a.smtpserver or ''}))
            else:
                print('unrecognized domain: ' + raw, file=sys.stderr)

    if a.limit:
        per, cut = {}, []
        for t, p in jobs:
            k = p['request']
            per[k] = per.get(k, 0) + 1
            if per[k] <= a.limit:
                cut.append((t, p))
        jobs = cut

    counts = {}
    for t, p in jobs:
        counts[p['request']] = counts.get(p['request'], 0) + 1
    print('to send: ' + ', '.join(f'{k}={v}' for k, v in sorted(counts.items())) if jobs else 'nothing to send')
    if not a.apply:
        for t, _ in jobs[:20]:
            print('  ' + t)
        if len(jobs) > 20:
            print(f'  ... and {len(jobs) - 20} more')
        print('dry run: nothing was sent, add --apply to send')
        return

    ok = fail = 0
    for t, p in jobs:
        try:
            r = api(url, token, p)
            if r.get('code') == 'success':
                ok += 1
            else:
                fail += 1
                print(f'FAILED {t}: {json.dumps(r, ensure_ascii=False)[:200]}', file=sys.stderr)
        except Exception as e:  # network or JSON error
            fail += 1
            print(f'FAILED {t}: {e}', file=sys.stderr)
        time.sleep(a.delay)
    print(f'done: ok={ok} failed={fail}')
    sys.exit(1 if fail else 0)


if __name__ == '__main__':
    main()
