#!/ usr/bin/env
# coding=utf-8
"""
author: b5mali4
Copyright (c) 2018
"""
import unittest
from peewee import *
from common import log
from exception.request_classification_exception import HeaderParseError


class VulnerabilityModelTestCase(unittest.TestCase):
    def testGetTasksUrlsVulnsNumByDays(self):
        """
        测试get_tasks_urls_vulns_num_by_days函数耗时操作，优化sql语句，查看具体的索引使用情况
        
        统计最近几天内的任务，url，漏洞数量
        SELECT * FROM task  WHERE DATE_SUB(CURDATE(), INTERVAL 10 DAY) <= created_time

        SELECT * FROM task INNER JOIN usertask ON usertask.task_id = task.id where usertask.user_id='1222'
        :return: 
        """
        import time
        from model.task import TaskService, Task
        from model.user_task import UserTask, UserTaskService
        from model.url import Url, UrlService
        from model.vulnerability import VulnerabilityService, Vulnerability
        from common.system_util import get_front_date
        from common.json_utils import dict_auto_add

        user_id = 1
        day_range = 100000

        tasks_num = {}
        urls_num = {}
        vulns_num = {}

        start_time = time.time()
        if user_id:

            task_total_num = TaskService.count(where=(Task.create_user_id == user_id))
        else:
            task_total_num = TaskService.count()
        url_total_num = UrlService.count()
        vuln_total_num = VulnerabilityService.count()

        print("总共消耗了{}".format(time.time() - start_time))

        if day_range:
            front_date = get_front_date(day_range=day_range)
            if user_id:
                matched_tasks = Task.select().join(UserTask, JOIN.INNER, on=(Task.id == UserTask.task_id)).where(
                    UserTask.user_id == user_id, Task.created_time >= front_date).execute()
            else:
                matched_tasks = Task.select().where(Task.created_time >= front_date).execute()
        else:
            matched_tasks = TaskService.get_fields_by_where()

        for matched_task in matched_tasks:
            create_day = str(matched_task.created_time).split(" ")[0]
            dict_auto_add(tasks_num, create_day)
            import logging
            logger = logging.getLogger('peewee')
            logger.addHandler(logging.StreamHandler())
            logger.setLevel(logging.DEBUG)
            tmp_url_num = UrlService.count(where=(Url.task_id == matched_task.id))
            tmp_vuln_num = VulnerabilityService.count(where=(Vulnerability.task_id == matched_task.id))
            dict_auto_add(urls_num, create_day, tmp_url_num)
            dict_auto_add(vulns_num, create_day, tmp_vuln_num)

        print("总共消耗了{}".format(time.time() - start_time))

        print({'task': len(matched_tasks), 'task_total_num': task_total_num, 'url': urls_num,
               'url_total_num': url_total_num,
               'vuln': vulns_num, 'vuln_total_num': vuln_total_num})

    def testGetVulnerabilitys_nlt(self):
        """
                根据任务id来组合获得 数量，等级，类型
        SELECT COUNT(*)	AS	`num`,SUM(if(`level` = 'high', 1, 0)) AS hight,SUM(if(`level` = 'low', 1, 0)) AS low,SUM(if(`level` = 'middle', 1, 0)) AS middle,	SUM(if(`type` = 'sql_inject', 1, 0)) AS sql_inject,	SUM(if(`type` = 'xss', 1, 0)) AS xss,	SUM(if(`type` = 'xxe', 1, 0)) AS xxe,	SUM(if(`type` = 'weak_pwd', 1, 0)) AS weak_pwd,	SUM(if(`type` = 'cmdect', 1, 0)) AS cmdect,	SUM(if(`type` = 'file_read', 1, 0)) AS file_read,	SUM(if(`type` = 'file_upload', 1, 0)) AS file_upload,	SUM(if(`type` = 'file_include', 1, 0)) AS file_include,	SUM(if(`type` = 'ddos', 1, 0)) AS ddos, SUM(if(`type` = 'csrf', 1, 0)) AS csrf,	SUM(if(`type` = 'cors', 1, 0)) AS cors, SUM(if(`type` = 'jsonp', 1, 0)) AS jsonp,	SUM(if(`type` = 'info_leak', 1, 0)) AS info_leak,	SUM(if(`type` = 'crlf', 1, 0)) AS crlf,	SUM(if(`type` = 'other', 1, 0)) AS other,	SUM(if(`type` = 'hidden_danger', 1, 0)) AS hidden_danger
FROM `vulnerability` AS vulnerability WHERE (vulnerability.`task_id` = 3)

vul_type = {'sql_inject': 0, 'xss': 0, 'xxe': 0, 'weak_pwd': 0, 'cmdect': 0, 'file_read': 0, 'file_upload': 0,
                    'file_include': 0, 'ddos': 0, 'csrf': 0,
                    'cors': 0, 'jsonp': 0, 'info_leak': 0, 'crlf': 0, 'other': 0, 'hidden_danger': 0}

        :return: 
        """
        from model.hunter_model import OrmModelJsonSerializer
        from model.vulnerability import Vulnerability, VulnerabilityService
        task_id = 2153
        import logging
        logger = logging.getLogger('peewee')
        logger.addHandler(logging.StreamHandler())
        logger.setLevel(logging.DEBUG)

        result = Vulnerability.select(fn.COUNT(Vulnerability.id).alias('num'),
                                      fn.SUM(fn.IF(Vulnerability.level == "high", 1, 0)).alias('high'),
                                      fn.SUM(fn.IF(Vulnerability.level == "middle", 1, 0)).alias('middle'),
                                      fn.SUM(fn.IF(Vulnerability.level == "low", 1, 0)).alias('low'),
                                      fn.SUM(fn.IF(Vulnerability.type == "sql_inject", 1, 0)).alias('sql_inject'),
                                      fn.SUM(fn.IF(Vulnerability.type == "xss", 1, 0)).alias('xss'),
                                      fn.SUM(fn.IF(Vulnerability.type == "xxe", 1, 0)).alias('xxe'),
                                      fn.SUM(fn.IF(Vulnerability.type == "weak_pwd", 1, 0)).alias('weak_pwd'),
                                      fn.SUM(fn.IF(Vulnerability.type == "cmdect", 1, 0)).alias('cmdect'),
                                      fn.SUM(fn.IF(Vulnerability.type == "file_read", 1, 0)).alias('file_read'),
                                      fn.SUM(fn.IF(Vulnerability.type == "file_upload", 1, 0)).alias('file_upload'),
                                      fn.SUM(fn.IF(Vulnerability.type == "file_include", 1, 0)).alias('file_include'),
                                      fn.SUM(fn.IF(Vulnerability.type == "ddos", 1, 0)).alias('ddos'),
                                      fn.SUM(fn.IF(Vulnerability.type == "csrf", 1, 0)).alias('csrf'),
                                      fn.SUM(fn.IF(Vulnerability.type == "cors", 1, 0)).alias('cors'),
                                      fn.SUM(fn.IF(Vulnerability.type == "jsonp", 1, 0)).alias('jsonp'),
                                      fn.SUM(fn.IF(Vulnerability.type == "info_leak", 1, 0)).alias('info_leak'),
                                      fn.SUM(fn.IF(Vulnerability.type == "crlf", 1, 0)).alias('crlf'),
                                      fn.SUM(fn.IF(Vulnerability.type == "other", 1, 0)).alias('other'),
                                      fn.SUM(fn.IF(Vulnerability.type == "hidden_danger", 1, 0)).alias(
                                          'hidden_danger')).where(
            Vulnerability.task_id == task_id).execute()[0]

        print(result.num)
        print(result.high)
        print(result.middle)
        print(result.low)
        print(OrmModelJsonSerializer.serializer(result))

    def test_get_tasks_urls_vulns_num_by_days(self):
        from model.url import Url
        from model.user import User
        from model.user_task import UserTask

        user_id = 1
        import logging
        logger = logging.getLogger('peewee')
        logger.addHandler(logging.StreamHandler())
        logger.setLevel(logging.DEBUG)
        url_total_num = Url.select(fn.COUNT(Url.id).alias('urls_total_num')).join(UserTask, JOIN.INNER, on=(
            UserTask.task_id == Url.task_id)).where(UserTask.user_id == user_id).execute()[0].urls_total_num


if __name__ == "__main__":
    unittest.main()
