# -*- coding: utf-8 -*-

import json

from ._internal_utils import status_code_check, response_status_check, convert_date_time
from datetime import datetime, date, timedelta
import logging
log = logging.getLogger(__name__)

import requests
import urlparse
import random


class Vulnerability(object):
    def __init__(self, console_url, session=None):
        self._console_url = console_url
        self._session = session
        self._host = urlparse.urlparse(console_url).hostname

    @property
    def session(self):
        return self._session

    @session.setter
    def session(self, value):
        self._session = value

    def import_vul(self, local_file, ip, api_key):
        """
        Import vulnerability info file
        :param local_file: vulnerability info file
        :return: None
        """
        uri = self._console_url + '/__internal/security/vulnerability/import'
        vuln_type = 1
        sequence = xrange(10000001, 99999999)
        header = {'x-api-key': api_key}

        with open(local_file) as f:
            vuln = json.load(f, encoding="UTF-8")

            for rec in vuln:
                params = {}
                params["ip"] = ip
                params["id"] = random.choice(sequence)
                params["type"] = vuln_type
                params["category"] = rec["category"]
                params["name"] = rec["name"]
                params["score"] = rec["score"]
                params["level"] = rec["level"]
                if "occurs" in rec:
                    params["occurs"] = rec["occurs"]
                
                if "description" in rec:
                    params["description"] = rec["description"]
                
                params["solution"] = rec["solution"]
        
                if "publish_time" in rec:
                    params["publish_time"] = rec["publish_time"]
                
                if "cncve_list" in rec:
                    params["cncve_list"] = rec["cncve_list"]
        
                if "cnnvd_list" in rec:
                    params["cnnvd_list"] = rec["cnnvd_list"]
                
                if "cnvd_list" in rec:
                    params["cnvd_list"] = rec["cnvd_list"]
                
                if "cve_list" in rec:
                    params["cve_list"] = rec["cve_list"]
                
                r = self._session.post(uri, headers=header, json=params)
                

    def vul_list(self, start_time=date.today(), end_time=date.today() + timedelta(days=1), filter_expression=None):
        """Get vulnerability list

            :param start_time: query start time
            :param end_time: query end time
            :return: vulnerability list
        """

        if isinstance(start_time, str):
            start_time = datetime.strptime(start_time, '%Y-%m-%d')
        if isinstance(end_time, str):
            end_time = datetime.strptime(end_time, '%Y-%m-%d')

        payload = {
            'paginate': False,
            'vulScene': {
                'startTime': convert_date_time(start_time),
                'endTime': convert_date_time(end_time),
            }
        }
        if filter_expression:
            payload['vulScene']['commonFilters'] = dict(FilterExpression=filter_expression)

        uri = self._console_url + '/api/node/vul/list'
        response = self._session.post(uri, json=payload)
        response_content = json.loads(response.content)
        status_code_check(response.status_code, 200)
        response_status_check(response_content['statusCode'], 0, response_content['messages'])
        return response_content['data']['list']

    def vul_detail(self, start_time=date.today(), end_time=date.today() + timedelta(days=1), filter=None):
        if isinstance(start_time, str):
            start_time = datetime.strptime(start_time, '%Y-%m-%d')
        if isinstance(end_time, str):
            end_time = datetime.strptime(end_time, '%Y-%m-%d')

        payload = {
            'paginate': False,
            'startTime': convert_date_time(start_time),
            'endTime': convert_date_time(end_time),
            'filter': filter
        }

        uri = self._console_url + '/api/node/vul/detail'
        response = self._session.post(uri, json=payload)
        response_content = json.loads(response.content)
        status_code_check(response.status_code, 200)
        response_status_check(response_content['statusCode'], 0, response_content['messages'])
        return response_content['data']['list']

    def vul_status(self, _id, confirm_status, reason):
        payload = {
            'confirm_status': confirm_status,
            'reason': reason,
            '_id': _id,
            '_index': "vulnerability_detail",
            '_type': "vulnerability"
        }
        uri = self._console_url + '/api/node/vul/status/update'
        response = self._session.post(uri, json=payload)
        response_content = json.loads(response.content)
        status_code_check(response.status_code, 200)
        response_status_check(response_content['statusCode'], 0, response_content['messages'])
        return response_content['data']

    def get_knowledge(self, id):
        uri = self._console_url + '/api/node/vul/knowledge/' + str(id)
        response = self._session.get(uri)
        response_content = json.loads(response.content)
        status_code_check(response.status_code, 200)
        response_status_check(response_content['statusCode'], 0, response_content['messages'])
        return response_content['data']

    def run_score(self, id='', ip=[]):
        uri = self._console_url + '/security/risk/run-score'
        payload = {
            "script_name": "score_all",
        }
        response = self._session.post(uri, json=payload)
        status_code_check(response.status_code, 200)

        score_group = {
            "script_name": "score_group",
            "params": {
                "values": ip,
                "scope": id
            }
        }
        response = self._session.post(uri, json=score_group)
        status_code_check(response.status_code, 200)