#!/usr/bin/python
# -*- coding: utf-8 -*-

# --------------------------------------------------------------------------- #
#                                                                             #
#    Plugin for iSida Jabber Bot                                              #
#                                                                             #
#    This program is free software: you can redistribute it and/or modify     #
#    it under the terms of the GNU General Public License as published by     #
#    the Free Software Foundation, either version 3 of the License, or        #
#    (at your option) any later version.                                      #
#                                                                             #
# --------------------------------------------------------------------------- #

import subprocess
import re
import os

# --------------------------------------------------------------------------- #
# CONFIGURATION                                                               #
# --------------------------------------------------------------------------- #

DNSRECON_OPTS = ['-n', '8.8.8.8,1.1.1.1', '--threads', '1', '--lifetime', '10', '--disable_check_recursion']
DNSRECON_TIMEOUT = 30

PROTOCOL_SRV = {
    'xmpp': ['xmpp-client', 'xmpp-server', 'xmpps-client', 'xmpps-server'],
    'call': ['turn', 'turns', 'stun', 'stuns'],
    'sip': ['sip', 'sips'],
    'mail': ['pop', 'pops', 'imap', 'imaps', 'smtp', 'smtps', 'submission'],
    'minecraft': ['minecraft'],
}

PROTOCOL_TXT = {
    'xmpp': ['_xmpp'],
    'dmarc': ['_dmarc'],
    'dkim': ['_domainkey'],
    'spf': ['v=spf1'],
    'sogo': ['_sogo'],
}

ALL_SRV = list(set([p for sublist in PROTOCOL_SRV.values() for p in sublist]))
ALL_TXT = list(set([p for sublist in PROTOCOL_TXT.values() for p in sublist]))

# --------------------------------------------------------------------------- #
# HELPERS                                                                     #
# --------------------------------------------------------------------------- #

def _is_valid_domain(domain):
    domain = domain.strip().lower()
    if not domain or len(domain) < 3 or domain.count('.') < 1:
        return False
    return bool(re.match(r'^[a-z0-9]([a-z0-9\-]*[a-z0-9])?(\.[a-z0-9]([a-z0-9\-]*[a-z0-9])?)*$', domain))

def _check_dnsrecon():
    for path in os.environ.get('PATH', '').split(':'):
        if os.path.exists(os.path.join(path, 'dnsrecon')):
            return True, None
    if os.path.exists('/usr/bin/apt') or os.path.exists('/usr/bin/apt-get'):
        return False, 'apt install dnsrecon'
    elif os.path.exists('/usr/bin/dnf'):
        return False, 'dnf install dnsrecon'
    elif os.path.exists('/usr/bin/yum'):
        return False, 'yum install dnsrecon'
    elif os.path.exists('/usr/bin/zypper'):
        return False, 'zypper install dnsrecon'
    elif os.path.exists('/usr/bin/pacman'):
        return False, 'pacman -S dnsrecon'
    else:
        return False, 'dnsrecon (unknown package manager)'

def _parse_filters(filter_str, proto_map):
    if filter_str.lower() == 'all':
        return list(set([p for sublist in proto_map.values() for p in sublist]))
    protocols = []
    for part in filter_str.lower().split(','):
        part = part.strip()
        if part in proto_map:
            protocols.extend(proto_map[part])
        else:
            protocols.append(part)
    return list(set(protocols))

def _run_dnsrecon(domain, record_type='srv'):
    cmd = ['dnsrecon', '-t', record_type, '-d', domain] + DNSRECON_OPTS
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=DNSRECON_TIMEOUT)
        lines = result.stdout.split('\n')
        records = []
        if record_type == 'srv':
            for line in lines:
                if 'SRV' in line.upper() and '_' in line:
                    records.append(line.strip())
        else:
            # TXT
            for line in lines:
                if 'TXT' in line.upper():
                    records.append(line.strip())
        return records, None
    except subprocess.TimeoutExpired:
        return None, L('Timeout: dnsrecon took too long (limit %s sec.)', '') % DNSRECON_TIMEOUT
    except FileNotFoundError:
        return None, L('dnsrecon not found', '')
    except Exception as e:
        return None, L('Error: %s', '') % str(e)

def _filter_srv(records, protocols):
    if not records:
        return []
    filtered = []
    for record in records:
        for proto in protocols:
            if '_%s._' % proto in record.lower():
                filtered.append(record)
                break
    return filtered

def _filter_txt(records, keywords):
    if not records:
        return []
    filtered = []
    for record in records:
        for kw in keywords:
            if kw.lower() in record.lower():
                filtered.append(record)
                break
    return filtered

def _deduplicate_srv(records):
    seen = set()
    unique = []
    for rec in records:
        # Убираем IP-адреса из строки для сравнения
        cleaned = re.sub(r'\b(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\b', '', rec)
        cleaned = re.sub(r'\b(?:[a-fA-F0-9]{1,4}:){7}[a-fA-F0-9]{1,4}\b|\b(?:[a-fA-F0-9]{1,4}:){1,7}:[a-fA-F0-9]{1,4}\b|\b::(?:[a-fA-F0-9]{1,4}:){0,6}[a-fA-F0-9]{1,4}\b', '', cleaned)
        cleaned = re.sub(r'\s+', ' ', cleaned).strip()
        if cleaned not in seen:
            seen.add(cleaned)
            unique.append(rec)
    return unique

def _format_srv(records):
    if not records:
        return L('No matching SRV records found', '')
    # Убираем дубликаты
    records = _deduplicate_srv(records)
    formatted = []
    for rec in records:
        ipv4 = re.findall(r'\b(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\b', rec)
        ipv6 = re.findall(r'\b(?:[a-fA-F0-9]{1,4}:){7}[a-fA-F0-9]{1,4}\b|\b(?:[a-fA-F0-9]{1,4}:){1,7}:[a-fA-F0-9]{1,4}\b|\b::(?:[a-fA-F0-9]{1,4}:){0,6}[a-fA-F0-9]{1,4}\b', rec)
        cleaned = rec
        for ip in ipv4 + ipv6:
            cleaned = cleaned.replace(ip, '')
        cleaned = re.sub(r'\s+', ' ', cleaned).strip()
        if ipv4 and ipv6:
            cleaned += ' | ipv4/ipv6'
        elif ipv4:
            cleaned += ' | ipv4'
        elif ipv6:
            cleaned += ' | ipv6'
        formatted.append(cleaned)
    return '\n'.join(formatted)

def _format_txt(records):
    if not records:
        return L('No matching TXT records found', '')
    # Убираем дубликаты
    seen = set()
    unique = []
    for rec in records:
        if rec not in seen:
            seen.add(rec)
            unique.append(rec)
    return '\n'.join(unique)

# --------------------------------------------------------------------------- #
# COMMAND HANDLERS                                                            #
# --------------------------------------------------------------------------- #

def dnsrecon_srv(type, jid, nick, text):
    text = text.strip()
    if not text or len(text.split()) < 2:
        send_msg(type, jid, nick, L('Usage: dnsrecon -srv <filter> <domain>\nExample: dnsrecon -srv xmpp,call ubunix.pro\nFilters: xmpp, call, sip, mail, minecraft, all', '%s/%s' % (jid, nick)))
        return
    parts = text.split()
    filter_str = parts[0]
    domain = ' '.join(parts[1:])
    if not _is_valid_domain(domain):
        send_msg(type, jid, nick, L('Invalid domain: %s', '%s/%s' % (jid, nick)) % domain)
        return
    has, install = _check_dnsrecon()
    if not has:
        send_msg(type, jid, nick, L('dnsrecon not installed.\nInstall: sudo %s', '%s/%s' % (jid, nick)) % install)
        return
    protocols = _parse_filters(filter_str, PROTOCOL_SRV)
    if not protocols:
        send_msg(type, jid, nick, L('No valid filters: %s', '%s/%s' % (jid, nick)) % filter_str)
        return
    records, error = _run_dnsrecon(domain, 'srv')
    if error:
        send_msg(type, jid, nick, error)
        return
    filtered = _filter_srv(records, protocols)
    if filtered:
        msg = L('SRV records for %s (filter: %s):\n%s', '%s/%s' % (jid, nick)) % (domain, filter_str, _format_srv(filtered))
    else:
        msg = L('No SRV records found for %s with filter: %s', '%s/%s' % (jid, nick)) % (domain, filter_str)
    send_msg(type, jid, nick, msg)

def dnsrecon_txt(type, jid, nick, text):
    text = text.strip()
    if not text or len(text.split()) < 2:
        send_msg(type, jid, nick, L('Usage: dnsrecon -txt <filter> <domain>\nExample: dnsrecon -txt xmpp,spf ubunix.pro\nFilters: xmpp, dmarc, dkim, spf, sogo, all', '%s/%s' % (jid, nick)))
        return
    parts = text.split()
    filter_str = parts[0]
    domain = ' '.join(parts[1:])
    if not _is_valid_domain(domain):
        send_msg(type, jid, nick, L('Invalid domain: %s', '%s/%s' % (jid, nick)) % domain)
        return
    has, install = _check_dnsrecon()
    if not has:
        send_msg(type, jid, nick, L('dnsrecon not installed.\nInstall: sudo %s', '%s/%s' % (jid, nick)) % install)
        return
    keywords = _parse_filters(filter_str, PROTOCOL_TXT)
    if not keywords:
        send_msg(type, jid, nick, L('No valid filters: %s', '%s/%s' % (jid, nick)) % filter_str)
        return
    records, error = _run_dnsrecon(domain, 'txt')
    if error:
        send_msg(type, jid, nick, error)
        return
    filtered = _filter_txt(records, keywords)
    if filtered:
        msg = L('TXT records for %s (filter: %s):\n%s', '%s/%s' % (jid, nick)) % (domain, filter_str, _format_txt(filtered))
    else:
        msg = L('No TXT records found for %s with filter: %s', '%s/%s' % (jid, nick)) % (domain, filter_str)
    send_msg(type, jid, nick, msg)

# --------------------------------------------------------------------------- #
# REGISTER COMMANDS                                                           #
# --------------------------------------------------------------------------- #

execute = [
    (6, 'dnsrecon -srv', dnsrecon_srv, 2, 'DNS SRV reconnaissance.\ndnsrecon -srv <filter> <domain>'),
    (6, 'dnsrecon -txt', dnsrecon_txt, 2, 'DNS TXT reconnaissance.\ndnsrecon -txt <filter> <domain>'),
]