hbyd_ueba/utils/dashboard_data_pg.py

#!/usr/bin/python
# encoding=utf-8
# author: tangwy

import json
import os, re
import codecs
import traceback
from collections import defaultdict

from isoc.utils.esUtil import EsUtil
from isoc.utils.dashboard_data_conversion import ip_summary_data_format, account_summary_data_format, \
    interface_summary_data_format, menu_summary_data_format, calculate_time_difference, summary_data_reqs_format
from dataInterface.functions import CFunction
from dataInterface.db.params import CPgSqlParam

from ext_logging import logger

TABLE_NAME = "ueba_logs"

DATA_TYPE = {
    "IP": 1,
    "ACCOUNT": 2,
    "INTERFACE": 3,
    "MENU": 4,
}


def pg_get_ip_group_data(startTime, endTime):
    """
        IP维度查询
        :param startTime: 开始时间,
        :param endTime: 结束时间,
    """
    result = []
    sql = """ select ip,  jobnum, sum(count) from {TABLE_NAME} 
                where logdate >= %s and logate < %s and data_type = %s 
                group by ip, jobnum""".format(TABLE_NAME=TABLE_NAME)
    res = CFunction.execute(CPgSqlParam(sql, params=(startTime, endTime, DATA_TYPE["IP"])))
    if res:
        for item in res:
            result.append({
                "ip": item[0],
                "jobnum": item[2],
                "count": item[3],
            })
    return result


def pg_get_account_group_data(startTime, endTime):
    """
        账号维度查询
        :param startTime: 开始时间,
        :param endTime: 结束时间,
    """
    result = []
    sql = """ select account,  jobnum, sum(count) from {TABLE_NAME} 
                where logdate >= %s and logate < %s and data_type = %s 
                group by account, jobnum""".format(TABLE_NAME=TABLE_NAME)
    res = CFunction.execute(CPgSqlParam(sql, params=(startTime, endTime, DATA_TYPE["ACCOUNT"])))
    if res:
        for item in res:
            result.append({
                "account": item[0],
                "jobnum": item[2],
                "count": item[3],
            })
    return result


def pg_get_interface_group_data(startTime, endTime):
    """
        接口维度查询
        :param startTime: 开始时间,
        :param endTime: 结束时间,
    """
    result = []
    sql = """ select interface, sip, jobnum, sum(count) from {TABLE_NAME} 
                where logdate >= %s and logate < %s and data_type = %s 
                group by interface, ip, jobnum""".format(TABLE_NAME=TABLE_NAME)
    res = CFunction.execute(CPgSqlParam(sql, params=(startTime, endTime, DATA_TYPE["INTERFACE"])))
    if res:
        for item in res:
            result.append({
                "interface": item[0],
                "ip": item[1],
                "jobnum": item[2],
                "count": item[3],
            })
    return result


def pg_get_menu_group_data(startTime, endTime):
    """
        菜单维度查询
        :param startTime: 开始时间,
        :param endTime: 结束时间,
    """
    result = []
    sql = """ select menu, sip, jobnum, sum(count) from {TABLE_NAME} 
                where logdate >= %s and logate < %s and data_type = %s 
                group by menu, ip, jobnum""".format(TABLE_NAME=TABLE_NAME)
    res = CFunction.execute(CPgSqlParam(sql, params=(startTime, endTime, DATA_TYPE["MENU"])))
    if res:
        for item in res:
            result.append({
                "menu": item[0],
                "ip": item[1],
                "jobnum": item[2],
                "count": item[3],
            })
    return result


def pg_get_previous_company_count(startTime, endTime, data_type):
    """
        账号维度查询菜请求次数
        :param startTime: 开始时间,
        :param endTime: 结束时间,
        :param data_type: 公司聚合类型 ACCOUNT or IP ,
    """
    result = defaultdict(int)
    if data_type in DATA_TYPE:
        data_type = DATA_TYPE[data_type]
        sql = """ select jobnum, sum(count) from {TABLE_NAME} 
                    where logdate >= %s and logate < %s and data_type = %s 
                    group by jobnum""".format(TABLE_NAME=TABLE_NAME)
        res = CFunction.execute(CPgSqlParam(sql, params=(startTime, endTime, data_type)))
        if res:
            for item in res:
                company = find_region_by_code(item[0], jobnum_region_dict)
                result[company] += item[1]
    return result


def pg_get_previous_interface_count(startTime, endTime):
    """
        接口维度查询请求总次数
         :param startTime: 开始时间,
        :param endTime: 结束时间,
    """
    result = defaultdict(int)
    sql = """ select interface, sum(count) from {TABLE_NAME} 
                where logdate >= %s and logate < %s and data_type = %s 
                group by interface""".format(TABLE_NAME=TABLE_NAME)
    res = CFunction.execute(CPgSqlParam(sql, params=(startTime, endTime, DATA_TYPE["INTERFACE"])))
    if res:
        for item in res:
            result[item[0]] += item[1]
    return result


def pg_get_previous_menu_count(startTime, endTime):
    """
        菜单维度查询请求总次数
         :param startTime: 开始时间,
        :param endTime: 结束时间,
    """
    result = defaultdict(int)
    sql = """ select menu, sum(count) from {TABLE_NAME} 
                where logdate >= %s and logate < %s and data_type = %s 
                group by menu""".format(TABLE_NAME=TABLE_NAME)
    res = CFunction.execute(CPgSqlParam(sql, params=(startTime, endTime, DATA_TYPE["MENU"])))
    if res:
        for item in res:
            result[item[0]] += item[1]
    return result


def entry(data_type, start, end):
    # 前一段开始时间
    previous_time = calculate_time_difference(start, end)

    try:
        data = {}
        if data_type == "1":
            ip_summary_data = pg_get_ip_group_data(start, end)
            data = ip_summary_data_format(ip_summary_data)

            previous_company_dict = pg_get_previous_company_count(previous_time, start, "IP")
            for d in data["summary"]["account"]:
                d["trend"] = round(
                    (d["req_frequency"] - previous_company_dict.get(d["company"], 0)) / previous_company_dict.get(
                        d["company"], 0), 4)

        if data_type == "2":
            account_summary_data = pg_get_account_group_data(start, end)
            data = account_summary_data_format(account_summary_data)

            previous_company_dict = pg_get_previous_company_count(previous_time, start, "ACCOUNT")
            for d in data["summary"]["account"]:
                d["trend"] = round(
                    (d["req_frequency"] - previous_company_dict.get(d["company"], 0)) / previous_company_dict.get(
                        d["company"], 0), 4)

        if data_type == "3":
            interface_summary_data = pg_get_interface_group_data(start, end)
            data = interface_summary_data_format(interface_summary_data)

            previous_interface_dict = pg_get_previous_interface_count(previous_time, start)
            for d in data["summary"]["account"]:
                d["trend"] = round(
                    (d["req_frequency"] - previous_interface_dict.get(d["company"], 0)) / previous_interface_dict.get(
                        d["company"], 0), 4)

        if data_type == "4":
            menu_summary_data = pg_get_menu_group_data(start, end)
            data = menu_summary_data_format(menu_summary_data)

            previous_menu_dict = pg_get_previous_menu_count(previous_time, start)
            for d in data["summary"]["account"]:
                d["trend"] = round(
                    (d["req_frequency"] - previous_menu_dict.get(d["company"], 0)) / previous_menu_dict.get(
                        d["company"], 0), 4)
        return data
    except Exception, e:
        logger.error("分析结构获取失败, err: {}, traceback: {}".format(str(e), traceback.format_exc()))
        raise e