import subprocess
import shlex
import random
import string
import psutil
import os
import re
import sys
import hashlib
from celery.utils.log import get_task_logger
import colorlog
import logging
import dns.resolver
from tld import get_tld
from datetime import datetime
try:
    from celery import current_task
except ImportError:
    current_task = None

from .conn import http_req, conn_db
from .http import get_title, get_headers
from .domain import check_domain_black, is_valid_domain, is_in_scope, is_in_scopes, is_valid_fuzz_domain
from .ip import is_vaild_ip_target, not_in_black_ips, get_ip_asn, get_ip_city, get_ip_type
from .arl import arl_domain, get_asset_domain_by_id
from .time import curr_date, time2date, curr_date_obj
from .url import rm_similar_url, get_hostname, normal_url, same_netloc, verify_cert, url_ext
from .cert import get_cert
from .arlupdate import arl_update
from .cdn import get_cdn_name_by_cname, get_cdn_name_by_ip
from .device import device_info
from .cron import check_cron, check_cron_interval
from .query_loader import load_query_plugins
import re

def load_file(path):
    with open(path, "r+", encoding="utf-8") as f:
        return f.readlines()


def exec_system(cmd, **kwargs):
    cmd = " ".join(cmd)
    timeout = 4 * 60 * 60

    if kwargs.get('timeout'):
        timeout = kwargs['timeout']
        kwargs.pop('timeout')

    completed = subprocess.run(shlex.split(cmd), timeout=timeout, check=False, close_fds=True, **kwargs)

    return completed


def check_output(cmd, **kwargs):
    cmd = " ".join(cmd)
    timeout = 4 * 60 * 60

    if kwargs.get('timeout'):
        timeout = kwargs.pop('timeout')

    if 'stdout' in kwargs:
        raise ValueError('stdout argument not allowed, it will be overridden.')

    output = subprocess.run(shlex.split(cmd), stdout=subprocess.PIPE, timeout=timeout, check=False,
               **kwargs).stdout
    return output


def random_choices(k=6):
    return ''.join(random.choices(string.ascii_lowercase + string.digits, k=k))


def gen_md5(s):
    return hashlib.md5(s.encode()).hexdigest()


class MongoSyslogHandler(logging.Handler):
    def emit(self, record):
        try:
            task_id = "global"
            # 尝试从 celery 运行上下文中提取 task_id
            if current_task:
                task_obj = current_task._get_current_object()
                if task_obj and hasattr(task_obj, 'arl_task_id'):
                    task_id = task_obj.arl_task_id
                elif current_task.request:
                    options = None
                    if current_task.request.args:
                        options = current_task.request.args[0]
                    elif current_task.request.kwargs and "options" in current_task.request.kwargs:
                        options = current_task.request.kwargs["options"]
                        
                    if isinstance(options, dict):
                        if "data" in options:
                            task_id = options["data"].get("task_id") or options["data"].get("job_id", "global")
                        elif "task_id" in options:
                            task_id = options.get("task_id", "global")
                        elif "job_id" in options:
                            task_id = options.get("job_id", "global")

            level = record.levelname.lower()
            if level == 'warn':
                level = 'warning'
            elif level in ['fatal', 'critical']:
                level = 'error'

            log_doc = {
                "task_id": task_id,
                "level": level,
                "title": getattr(record, 'funcName', '系统运行'),
                "message": str(record.getMessage()),
                "create_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            }
            conn_db('syslog').insert_one(log_doc)
        except Exception:
            pass # 捕获异常防止拖垮主业务进程

def init_logger():
    handler = colorlog.StreamHandler()
    handler.setFormatter(colorlog.ColoredFormatter(
        fmt = '%(log_color)s[%(asctime)s] [%(levelname)s] '
              '[%(threadName)s] [%(filename)s:%(lineno)d] %(message)s', datefmt = "%Y-%m-%d %H:%M:%S"))

    logger = colorlog.getLogger('arlv2')

    logger.setLevel(logging.INFO)
    logger.addHandler(handler)
    
    # 💥 添加我们自定义的 MongoDB 拦截器
    mongo_handler = MongoSyslogHandler()
    mongo_handler.setLevel(logging.INFO)
    logger.addHandler(mongo_handler)
    
    logger.propagate = False


def get_logger():
    """
    根据类型创建日志记录器。
    :return: 日志记录器
    """
    if 'celery' in sys.argv[0]: # 内存字符串的物理匹配
        task_logger = get_task_logger(__name__)
        # 确保 task_logger 也挂载了我们的 MongoDB 拦截器
        has_mongo_handler = any(isinstance(h, MongoSyslogHandler) for h in task_logger.handlers)
        if not has_mongo_handler:
            mongo_handler = MongoSyslogHandler()
            mongo_handler.setLevel(logging.INFO)
            task_logger.addHandler(mongo_handler)
        return task_logger

    logger = logging.getLogger('arlv2')
    if not logger.handlers:
        init_logger()

    return logger


def get_ip(domain, log_flag=True):
    domain = domain.strip()
    logger = get_logger()
    ips = []
    try:
        answers = dns.resolver.resolve(domain, 'A')
        for rdata in answers:
            if rdata.address == '0.0.0.1':
                continue
            ips.append(rdata.address)
    except dns.resolver.NXDOMAIN as e:
        if log_flag:
            logger.info("{} {}".format(domain, e))

    except Exception as e:
        if log_flag:
            logger.warning("{} {}".format(domain, e))

    return ips


def get_cname(domain, log_flag=True):
    logger = get_logger()
    cnames = []
    try:
        answers = dns.resolver.resolve(domain, 'CNAME')
        for rdata in answers:
            cnames.append(str(rdata.target).strip(".").lower())
    except dns.resolver.NoAnswer as e:
        if log_flag:
            logger.debug(e)
    except Exception as e:
        logger.warning("{} {}".format(domain, e))

    return cnames


def domain_parsed(domain, fail_silently=True):
    """
    将域名拆分为主域名、域名、子域名
    :param domain:
    :param fail_silently:
    :return:
    """
    domain = domain.strip()
    try:
        res = get_tld(domain, fix_protocol=True,  as_object=True)
        item = {
            "subdomain": res.subdomain,
            "domain":res.domain,
            "fld": res.fld
        }
        return item
    except Exception as e:
        if not fail_silently:
            raise e


def get_fld(d):
    """获取域名的主域"""
    res = domain_parsed(d)
    if res:
        return res["fld"]


def gen_filename(site):
    filename = site.replace('://', '_')

    return re.sub(r'[^\w\-_\\. ]', '_', filename)


def build_ret(error, data):
    if isinstance(error, str):
        error = {
            "message": error,
            "code": 999,
        }

    ret = {}
    ret.update(error)
    ret["data"] = data
    msg = error["message"]

    if error["code"] != 200:
        for k in data:
            if k.endswith("id"):
                continue
            if not data[k]:
                continue
            if isinstance(data[k], str):
                msg += " {}:{}".format(k, data[k])

    ret["message"] = msg
    return ret


def kill_child_process(pid):
    logger = get_logger()
    parent = psutil.Process(pid)
    for child in parent.children(recursive=True):
        logger.info("kill child_process {}".format(child))
        child.kill()


def exit_gracefully(signum, frame):
    logger = get_logger()
    logger.info('Receive signal {} frame {}'.format(signum, frame))
    pid = os.getpid()
    kill_child_process(pid)
    parent = psutil.Process(pid)
    logger.info("kill self {}".format(parent))
    parent.kill()


def truncate_string(s):
    if len(s) > 30:
        truncated_string = s[:30]
        return truncated_string + "..."
    else:
        return s


def is_valid_exclude_ports(exclude_ports):
    """
    检查 nmap 中的排除端口范围是否合法
    """
    port_pattern = r'(\d+(-\d+)?,?)+'

    match = re.fullmatch(port_pattern, exclude_ports)

    if match:
        parts = exclude_ports.split(',')
        for part in parts:
            if '-' in part:
                start, end = map(int, part.split('-'))
                if start > end or not (0 <= start <= 65535) or not (0 <= end <= 65535):
                    return False
            else:
                if not (0 <= int(part) <= 65535):
                    return False
        return True
    else:
        return False


from .user import user_login, user_login_header, auth, user_logout, change_pass
from .push import message_push
from .fingerprint import parse_human_rule, transform_rule_map


def clean_task_data(task_id):
    """
    清理指定任务运行过程中产生的各表数据，用于任务无副作用的重新执行
    """
    logger = get_logger()
    logger.info(f"Cleaning existing data for task {task_id}")
    table_list = [
        "cert", "domain", "fileleak", "ip", "service",
        "site", "url", "vuln", "cip", "npoc_service", 
        "wih", "nuclei_result", "stat_finger"
    ]
    for table_name in table_list:
        try:
            conn_db(table_name).delete_many({'task_id': task_id})
        except Exception as e:
            logger.error(f"Error cleaning {table_name} for task {task_id}: {e}")
