# 任务管理接口
import json
import uuid

from apiflask import APIBlueprint, Schema
from apiflask.fields import String, Integer, IP, DateTime, List, Nested
from apiflask.validators import OneOf

from exts import db
from model import TaskPolicy
from .agentcomm import deliver_task
from .policy import *
from .util import da, error

bp = APIBlueprint("任务管理接口集合", __name__, url_prefix="/task")

# 数据库列与返回值的键对应关系
task_response_map = {
    "TASK_ID": "id",
    "TARGET_IP": "target",
    "TASK_NAME": "name",
    "AGENT_ID": "agent",
    "TARGET_DOMAIN": "target_domain",
    "TARGET_RR": "target_rr",
    "POLICY": "policy",
    "CREATE_TIME": "create_time",
    "STATUS": "status",
}


class Task(Schema):
    id = String()
    target = String()
    name = String()
    agent = String()
    target_domain = String()
    target_rr = String()
    policy = String()
    create_time = DateTime()
    status = String(validate=OneOf(["working", "stop", "finish"]))


# 执行输出日志
class TaskLog(Schema):
    # 输出时间
    time = DateTime()
    # 输出来源IP
    ip = String()
    # 该策略的目标IP
    targetip = String()
    # 日志等级
    level = String(validate=OneOf(["INFO", "WARNING", "ERROR"]))
    # 日志内容
    info = String()


# 任务状态时间轴信息
class TaskState(Schema):
    # 策略开始执行时间
    start_time = DateTime()
    # 策略名称
    policy_name = String()
    # 策略参数
    policy_param = String()
    # 执行策略编号
    policy_id = String()
    # 效果评估
    policy_status = String()


# 任务创建参数逻辑合法性检查
def valid_task_info(task_param: dict):
    if task_param["policy"] == "sjqp":
        if task_param["target_rr"] == "" or task_param["tdomain"] == "":
            return "数据欺骗缺乏目标域名或注入参数"
    return None


# 创建任务接口
@bp.post("/create")
@bp.doc("任务创建接口", "部分字段值的映射关系：</br>"
        + "policy 期望策略，可选参数范围及对应含义为：，auto-自动，ddos-拒绝服务，sjqp-数据欺骗</br>"
        + "scan 状态感知方式，可选参数范围及对应含义为：auto-自动，icmp-icmp/v6时延，tcp-tcp时延，dns-dns时延，record-记录正确性验证")
@bp.input({
    # 任务名称
    "name": String(),
    # 目标IP
    "target": IP(),
    # 执行代理
    "agent": String(),
    # 目标域名
    "target_domain": String(),
    # 期望注入记录
    "target_rr": String(),
    # 期望策略
    "policy": String(validate=OneOf(["auto", "ddos", "sjqp"])),
    # 状态感知方式
    "scan": String(validate=OneOf(["auto", "icmp", "dns", "tcp", "record"])),
    # 策略切换时限,单位分钟
    "policy_time": Integer(),
    # 任务执行时限，单位分钟
    "run_time": Integer(),
    # 运行配置
    "run_flag": String(validate=OneOf(["now", "man"]))
}, example={'name': "test_task", 'target': "1.2.3.4", 'agent': "8a9ces", 'target_domain': "www.google.com",
            'target_rr': "NS ns.ourattack.com",
            'policy': "auto", 'scan': "auto",
            'policy_time': 60, "run_time": 600, 'run_flag': "now"})
@bp.output({
    "code": Integer(),
    "msg": String()
})
# TODO：TARGET处理
def make_task(json_data):
    # 请求参数转储到变量
    name = json_data["name"]
    target = str(json_data["target"])
    agent_id = json_data["agent"]
    tdomain = json_data["target_domain"]
    target_rr = json_data["target_rr"]
    policy_type = json_data["policy"]
    scan = json_data["scan"]
    policy_time = json_data["policy_time"]
    run_flag = json_data["run_flag"]
    run_time = json_data["run_time"]

    # 任务编号生成
    t_id = str(uuid.uuid1())

    # 任务信息生成

    task_info = {
        "t_id": t_id,
        "name": name,
        "target": target,
        "ag_id": agent_id,
        "p_delay": policy_time,
        "t_delay": run_time,
        "scan": scan,
        "status": "working" if run_flag == "now" else "stop",
        "policy": policy_type,
        "tdomain": tdomain,
        "target_rr": target_rr,
        "target_rtype": ""
    }

    # 输入参数检查处理
    warn = valid_task_info(task_info)
    if warn is not None:
        return {"code": 400, "msg": warn}
    if task_info["target_rr"] != "":
        rr = task_info["target_rr"].split(" ")
        task_info["target_rtype"], task_info["target_rr"] = rr[0], rr[1]

    # 添加表名
    task_info["tab"] = MYSQL_TAB_TASK
    # 若目标为IP地址
    if is_ipaddress(target) is not None:
        task_info["TARGET"] = "TARGET_IP"
    # 若目标为域名，视为DoH服务
    else:
        task_info["TARGET"] = "TARGET_DOH"

    # 添加状态探测节点，用于状态感知
    scan_node_sql = """
    SELECT AGENT_ID,IPADDR
    FROM %s
    WHERE STATUS='1' AND AGENT_TYPE='ztgz'
    LIMIT 50;
    """ % MYSQL_TAB_AGENT
    da.cursor.execute(scan_node_sql)
    # 所有在线的支持状态感知的代理
    all_nodes = da.cursor.fetchall()
    # 随机选择10个
    selected_nodes = random.choices(all_nodes, k=10 if len(all_nodes) > 10 else len(all_nodes))
    # 将选择的代理节点ID和对应IPv4地址存储为字典，键为代理ID，值为代理的IPv4地址
    selected_nodes_info = {'{}'.format(n['AGENT_ID']): '{}'.format(str(n['IPADDR']).split("|")[0]) for n in
                           selected_nodes}
    task_info["nodes"] = json.dumps(selected_nodes_info)

    # sql语句组合
    sql = """
    INSERT INTO %(tab)s (
    TASK_ID,
    TASK_NAME, 
    AGENT_ID,  
    STATUS, 
    POLICY_DELAY,
    TASK_DELAY,
    TARGET_SCAN,
    TARGET_DOMAIN,
    TARGET_RTYPE,
    TARGET_RR,
    %(TARGET)s,
    SCAN_AGENT_ID_LIST)
    VALUES (
    '%(t_id)s',
    '%(name)s',
    '%(ag_id)s',
    '%(status)s',
    '%(p_delay)s',
    '%(t_delay)s',
    '%(scan)s',
    '%(tdomain)s',
    '%(target_rtype)s',
    '%(target_rr)s',
    '%(target)s',
    '%(nodes)s'
    );
    """ % task_info
    da.conn.ping(reconnect=True)
    da.cursor.execute(sql)
    da.conn.commit()

    # 任务策略初始化
    tp_id, p_exe, p_param = init_task_policy(policy_type, target, t_id)
    update_policy_sql = """
    UPDATE %s
    SET POLICY='%s'
    WHERE TASK_ID='%s';
    """ % (MYSQL_TAB_TASK, tp_id, t_id)
    da.cursor.execute(update_policy_sql)
    da.conn.commit()

    # 根据run_flag判断是否立刻执行
    if run_flag == "now":
        err = deliver_task(agent_id, p_exe, p_param)
        if err is not None:
            error(str(err))
            return {"code": 500, "msg": str(err)}
    return {"code": 200, "msg": "ok"}


opsmap = {"start": "开始", "stop": "暂停", "cancel": "停止"}
statemap = {"working": "执行中", "stop": "暂停", "finish": "结束"}
# 合法逻辑映射，例如不能让一个暂停的任务开始，或是让一个正在运行的任务停止
ops_state = {"start": ["stop"], "stop": ["working"], "cancel": ["working", "stop"]}
# 操作与状态结果的映射关系
ops_state_result = {"start": "working", "stop": "stop", "cancel": "finish"}


# 操作任务开始停止控制接口
@bp.post("/ops")
@bp.doc("任务操作接口")
@bp.input({
    "taskid": String(required=True),
    "ops": String(required=True, validate=OneOf(["start", "stop", "cancel"]))
})
@bp.output({
    "code": Integer(),
    "msg": String(),
})
def ops_task(json_data):
    ops = json_data["ops"]
    # 操作中文名
    ops_cn = opsmap[json_data["ops"]]
    id = json_data["taskid"]

    # 查询目标任务状态
    sql = """
    SELECT AGENT_ID,STATUS 
    FROM %s
    WHERE TASK_ID='%s';
    """ % (MYSQL_TAB_TASK, id)
    da.cursor.execute(sql)
    data = da.cursor.fetchall()
    if len(data) == 0:
        return {"code": 400, "msg": "任务" + id + "不存在！！"}
    # 判断操作是否是逻辑上可执行的
    if data[0]["STATUS"] in ops_state[ops]:
        # 向对应代理下发控制指令
        err = deliver_task(data[0]["AGENT_ID"], "ops", ops)
        if err is not None:
            return {"code": 500, "msg": str(err)}
        task_update_status_sql = """
        UPDATE %s
        SET STATUS='%s'
        WHERE TASK_ID='%s';""" % (MYSQL_TAB_TASK, ops_state_result[ops], id)
        da.cursor.execute(task_update_status_sql)
        da.conn.commit()
        return {"code": 200, "msg": "任务" + json_data["taskid"] + "已" + ops_cn}
    else:
        return {"code": 400, "msg": "无法对一个" + statemap[data[0]["STATUS"]] + "的任务执行" + ops_cn + "操作"}


# 查询任务列表接口
@bp.get("/")
@bp.doc("任务列表信息获取接口")
@bp.input({
    "page": Integer(load_default=1),
    "per_page": Integer(load_default=10)
}, location="query")
@bp.output({
    "code": Integer(),
    "data": List(Nested(Task())),
    "total": Integer()
})
# TODO：DoH适配
def tasks_state(query_data):
    per_page = query_data["per_page"]
    page = query_data["page"]

    # 任务列表
    task_list = []
    res = da.get_data(data_type="task",
                      offset=(page - 1) * per_page, limit=per_page)
    res_count = da.count_data(data_type="task")
    for r in res:
        task = {}
        for key, value in r.items():
            try:
                task[task_response_map[key]] = value
            except KeyError:
                continue
        task_list.append(task)
    return {"code": 200, "data": task_list, "total": res_count}


# 任务详情接口
@bp.get("/detail")
@bp.doc("任务执行状态时间轴信息获取接口")
@bp.input({
    "taskid": String(required=True),
}, location="query")
@bp.output({
    "code": Integer(),
    "data": List(Nested(TaskState()))
})
def task_info(query_data):
    taskid = query_data["taskid"]

    # 任务状态列表
    sql = """
    SELECT tp.TP_TIME as tp_time, p.P_NAME as p_name, tp.POLICY_PARAM as tp_param,tp.TP_ID as tp_id
    FROM %s as tp,%s as p
    WHERE tp.POLICY=p.P_ID AND tp.FOR_TASK='%s'
    ORDER BY TP_TIME;
    """ % (MYSQL_TAB_TASKPOLICY, MYSQL_TAB_POLICY, taskid)
    da.cursor.execute(sql)
    data = da.cursor.fetchall()
    task_state_list = []
    # 判断是否不存在该任务
    if len(data) == 0:
        return {"code": 404, "data": []}
    else:
        # 过往记录
        for d in data:
            task_state_list.append({
                "start_time": d["tp_time"],
                "policy_name": d["p_name"],
                "policy_param": d["tp_param"],
                "policy_id": d["tp_id"],
                "policy_status": "无效；原因为：超时未成功"
            })
        # 最后一项为当前正在执行的策略
        task_state_list[-1]["policy_status"] = "评估中"
    return {"code": 200, "data": task_state_list}


@bp.get("/tp")
@bp.doc("任务策略执行日志获取接口")
@bp.input({
    "id": Integer(required=True),
    "per_page": Integer(load_default=10),
    "page": Integer(load_default=1)
}, location="query")
@bp.output({
    "code": Integer(),
    "data": List(Nested(TaskLog())),
    "total": Integer()
})
def taskpolicy_log(query_data):
    id = query_data["id"]
    per_page = query_data["per_page"]
    page = query_data["page"]
    task_policy = db.session.get(TaskPolicy, id)
    if task_policy is None:
        return {"code": 404}

    logs = task_policy.task_logs
    # 分页
    total = len(logs)
    start = (page - 1) * per_page
    end = start + per_page
    paginated_logs = logs[start:end]

    log_data = [{
        "time": log.created_time,
        "targetip": log.task_policy.task.target_ip,
        "ip": log.task_policy.task.agent.ipaddr,
        "level": log.tlog_level,
        "info": log.tlog_info
    } for log in paginated_logs]

    return {"code": 200,
            "data": log_data,
            "total": total}


@bp.doc("任务删除接口", "输入参数说明：</br>"
        + "id: 任务编号")
@bp.post("/del")
@bp.input({
    "id": String(required=True)
})
@bp.output({
    "code": Integer(),
    "msg": String()
})
def del_task(json_data):
    taskid = json_data["id"]
    # 首先删除所有有关联的执行策略
    policysql = """
    DELETE
    FROM %s
    WHERE FOR_TASK='%s'
    """ % (MYSQL_TAB_TASKPOLICY, taskid)
    try:
        da.cursor.execute(policysql)
        da.conn.commit()
    except Exception as e:
        error(str(e))
        return {"code": 500, "msg": str(e)}
    # 删除任务记录
    sql = """
    DELETE 
    FROM %s 
    WHERE TASK_ID='%s';
    """ % (MYSQL_TAB_TASK, taskid)
    try:
        da.cursor.execute(sql)
        da.conn.commit()
        if da.cursor.rowcount == 0:
            return {"code": "404", "msg": "该任务id未找到！"}
    except Exception as e:
        error(str(e))
        return {"code": 500, "msg": str(e)}
    return {"code": 200, "msg": "ok"}