from goby_sdk.V1.lib.base_page import BasePage


class VulnManagementPage(BasePage):

    def get_pocs(self, task_id, query, reload_pocs=False, order_by='vul_nums', order='desc', page=1, page_size=20):
        """
        获取POC列表

        :param task_id: 任务ID
        :type task_id: str
        :param query: 查询条件
        :type query: str
        :param reload_pocs: 是否重新加载POC，默认为False
        :type reload_pocs: bool
        :param order_by: 排序字段，默认为"vul_nums"
        :type order_by: str
        :param order: 排序方式，默认为"desc"
        :type order: str
        :param page: 页码，默认为1
        :type page: int
        :param page_size: 每页数量，默认为20
        :type page_size: int
        :return: 返回POC列表的JSON对象
        :rtype: dict
        """
        options = {
            "reloadPocs": reload_pocs,
            "order": {
                "vul_nums": "desc",
                "level": "desc",
                "host_nums": "desc"
            },
            "page": {
                "page": 1,
                "size": 20
            }
        }

        if order_by:
            options['order'][order_by] = order

        if page or page_size:
            options['page']['page'] = page
            options['page']['size'] = page_size

        data = {
            "taskId": task_id,
            "query": query,
            "options": options
        }

        return self._send_request("api/v1/getPocs", data)

    def get_failed_pocs(self):
        """
        获取失败的POC列表

        :return: 返回失败的POC列表的JSON对象
        :rtype: dict
        """
        return self._send_get_request("api/v1/getFailedPocs")

    def get_vuln_statistics(self, task_id):
        """
        获取任务的漏洞统计数据

        :param task_id: 任务ID
        :type task_id: str
        :return: 返回漏洞统计数据的JSON对象
        :rtype: dict
        """
        data = {"taskId": task_id}
        return self._send_request("api/v1/vulnerabilityStatisticsData", data)

    def get_vulnerability_search(self, task_id, page=1, size=20):
        """
        获取任务的漏洞搜索结果

        :param task_id: 任务ID
        :type task_id: str
        :param page: 搜索结果的页码，默认为1
        :type page: int
        :param size: 每页的搜索结果数量，默认为20
        :type size: int
        :return: 返回漏洞搜索结果的JSON对象
        :rtype: dict
        """
        data = {
            "taskId": task_id,
            "type": "vulnerability",
            "query": f"taskId=\"{task_id}\"",
            "options": {
                "order": {"level": "desc", "nums": "desc"},
                "page": {"page": page, "size": size}
            }
        }
        return self._send_request("api/v1/vulnerabilitySearch", data)

    def get_poc_info(self, vul_name):
        # 视乎还有个hostinfo参数，未添加
        """
        获取POC信息

        :param vul_name: 漏洞名称
        :type vul_name: str
        :return: 返回POC信息的JSON对象
        :rtype: dict
        """
        data = {"vulname": vul_name}
        return self._send_request("api/v1/getPOCInfo", data)
