#!/usr/bin/env python3

# *******************************************************
# Copyright (c) VMware, Inc. 2021. All Rights Reserved.
# SPDX-License-Identifier: MIT
# *******************************************************
# *
# * DISCLAIMER. THIS PROGRAM IS PROVIDED TO YOU "AS IS" WITHOUT
# * WARRANTIES OR CONDITIONS OF ANY KIND, WHETHER ORAL OR WRITTEN,
# * EXPRESS OR IMPLIED. THE AUTHOR SPECIFICALLY DISCLAIMS ANY IMPLIED
# * WARRANTIES OR CONDITIONS OF MERCHANTABILITY, SATISFACTORY QUALITY,
# * NON-INFRINGEMENT AND FITNESS FOR A PARTICULAR PURPOSE.

"""Model and Query Classes for Vulnerability Assessment API"""

from cbc_sdk.base import (NewBaseModel, BaseQuery, QueryBuilder, QueryBuilderSupportMixin,
                          IterableQueryMixin, AsyncQueryMixin)
from cbc_sdk.errors import ApiError
import logging
import time

log = logging.getLogger(__name__)

""" Vulnerability models """


class OrganizationalVulnerability(NewBaseModel):
    """Represents a organizational vulnerability"""
    urlobject = '/vulnerability/assessment/api/v1/orgs/{0}'
    swagger_meta_file = "workload/models/paginatedDeviceVulnerabilityList.yaml"
    url_additional = ""

    def __init__(self, cb, model_unique_id=None, initial_data=None):
        """
        Initialize the OrganizationalVulnerability object.

        Args:
            cb (BaseAPI): Reference to API object used to communicate with the server.
            model_unique_id (int): Not used by this class
            initial_data (dict): dictionary of the data
        """
        super(OrganizationalVulnerability, self).__init__(cb, model_unique_id, initial_data)
        self._full_init = True

    @classmethod
    def _query_implementation(cls, cb, **kwargs):
        """
        Returns the appropriate query object for this object type.

        Args:
            cb (BaseAPI): Reference to API object used to communicate with the server.
            **kwargs (dict): Not used, retained for compatibility.

        Returns:
            VulnerabilityQuery: The query object for this vulnerability type.
        """
        return VulnerabilityQuery(cls, cb)


class VulnerabilitySummary(NewBaseModel):
    """Represents a vulnerability summary."""
    urlobject = "/vulnerability/assessment/api/v1/orgs/{0}"
    swagger_meta_file = "workload/models/organizationVulnerabilitySummary.yaml"
    url_additional = "/summary"

    VALID_SEVERITY = ["CRITICAL", "IMPORTANT", "MODERATE", "LOW"]

    def __init__(self, cb, model_unique_id=None, initial_data=None):
        """
        Initialize VulnerabilitySummary object

        Args:
            cb (BaseAPI): Reference to API object used to communicate with the server.
            model_unique_id (int): Not used by this class
            initial_data (dict): dictionary of the data
        """
        super(VulnerabilitySummary, self).__init__(cb, model_unique_id, initial_data)
        self._full_init = True

    @classmethod
    def _query_implementation(cls, cb, **kwargs):
        """
        Returns the appropriate query object for this object type.

        Args:
            cb (BaseAPI): Reference to API object used to communicate with
            the server.
            **kwargs (dict): Not used, retained for compatibility.

        Returns:
            VulnerabilityQuery: The query object for this vulnerability type.
        """
        return VulnerabilityQuery(cls, cb)

    @classmethod
    def get_org_vulnerability_summary(cls, cb, severity=None, vcenter_id=None):
        """
        Returns vulnerability summary at the organization level

        Args:
            cb (BaseAPI): Reference to API object used to communicate with the server.
            severity (str): (optional) filters the vulnerability sumary per severity
            (CRITICAL, IMPORTANT, MODERATE, LOW)
            vcenter_id (str): (optional) filters the vulenerability per vcenter id
        Returns:
            dict: monitored_assets and severity summary
        """
        query_params = {}
        if vcenter_id:
            url = "/vcenters/{}/vulnerabilities/summary".format(vcenter_id)
        else:
            url = "/vulnerabilities/summary"

        if severity and severity not in VulnerabilitySummary.VALID_SEVERITY:
            raise ApiError("Invalid severity")
        elif severity:
            query_params["severity"] = severity

        req_url = OrganizationalVulnerability.urlobject.format(
            cb.credentials.org_key) + url
        return cb.get_object(req_url, query_params)


class Vulnerability(NewBaseModel):
    """Represents a vulnerability"""
    urlobject = '/vulnerability/assessment/api/v1'
    urlobject_single = '/vulnerability/assessment/api/v1/vulnerabilities/{}'
    swagger_meta_file = "workload/models/vulnerability.yaml"
    url_additional = ""
    primary_key = "cve_id"

    def __init__(self, cb, model_unique_id, initial_data=None):
        """
        Initialize the Vulnerability object.

        Args:
            cb (BaseAPI): Reference to API object used to communicate with the server.
            model_unique_id (str): ID of the vulnerability represented.
            initial_data (dict): Initial data used to populate the alert.
        """
        super(Vulnerability, self).__init__(cb, model_unique_id, initial_data)
        if model_unique_id is not None and initial_data is None:
            self._refresh()
        self._full_init = True
        self._os_product_id = None

    def _refresh(self):
        """
        Rereads the object data from the server.

        Returns:
            bool: True if refresh was successful, False if not.
        """
        resp = self._cb.get_object(self._build_api_request_uri())
        resp['cve_id'] = self._model_unique_id
        self._info = resp
        self._last_refresh_time = time.time()
        return True

    def _build_api_request_uri(self, http_method="GET"):
        """
        Build the unique URL used to make requests for this object.

        Args:
            http_method (str): Not used; retained for compatibility.

        Returns:
            str: The URL used to make requests for this object.
        """
        return self.urlobject_single.format(self._model_unique_id)

    def affected_assets(self, os_product_id):
        """Returns a list of Vulnerability objects associated with device.

        Args;
            os_product_id (str) operating system product ID
        Returns:
            AffectedAssetQuery
        """
        if os_product_id:
            self._os_product_id = os_product_id
            return AffectedAssetQuery(self, self._cb)
        else:
            raise ApiError("The os product id is required parameter")

    @classmethod
    def _query_implementation(cls, cb, **kwargs):
        """
        Returns the appropriate query object for this object type.

        Args:
            cb (BaseAPI): Reference to API object used to communicate with the server.
            **kwargs (dict): Not used, retained for compatibility.

        Returns:
            AffectedAssetQuery: The query object
        """
        return AffectedAssetQuery(cls, cb)


class DeviceVulnerability(NewBaseModel):
    """Represents devices affected by a vulnerability"""
    urlobject = '/vulnerability/assessment/api/v1/orgs/{}/'
    swagger_meta_file = "workload/models/vulnerability.yaml"
    url_additional = "devices/{}/vulnerabilities/summary"

    VALID_CATEGORY = ["OS", "APP"]

    def __init__(self, cb, model_unique_id=None, initial_data=None):
        """
        Initialize DeviceVulnerability

        Args:
            cb (BaseAPI): Reference to API object used to communicate with the server.
            model_unique_id (str): ID of the vulnerability represented.
            initial_data (dict): Initial data used to populate the alert.
        """
        super(DeviceVulnerability, self).__init__(cb, model_unique_id, initial_data)
        self._full_init = True

    @classmethod
    def _query_implementation(cls, cb, **kwargs):
        """
        Returns the appropriate query object for this object type.

        Args:
            cb (BaseAPI): Reference to API object used to communicate with the server.
            **kwargs (dict): Not used, retained for compatibility.

        Returns:
            DeviceVulnerabilityQuery: The query object
        """
        return DeviceVulnerabilityQuery(cls, cb)

    @classmethod
    def get_vulnerability_summary_per_device(cls, cb, device_id, category=None, vcenter_id=None):
        """
        Returns vulnerability summary at the device level

        Args:
            cb (BaseAPI): Reference to API object used to communicate with the server.
            device_id (int) the device id for which the vulnerability data is returned
            category (str) (optional) category for which the vulnerability data is filtered (OS, APP)
            vcenter_id (str) (optional) ID of the vcenter for which the vulnerability data is returned

        Returns:
            dictinary: summary for vulnerabilities per device
        """
        query_params = {}
        url = ''
        if vcenter_id:
            url = "vcenters/{}/".format(vcenter_id)

        if category and category not in DeviceVulnerability.VALID_CATEGORY:
            raise ApiError("Invalid category provided")
        elif category:
            query_params["category"] = category

        req_url = DeviceVulnerability.urlobject.format(
            cb.credentials.org_key) + url + DeviceVulnerability.url_additional.format(device_id)
        return cb.get_object(req_url, query_params)


class VulnerabilityQuery(BaseQuery, QueryBuilderSupportMixin,
                         IterableQueryMixin, AsyncQueryMixin):
    """Represents a query that is used to locate Vulnerabiltity objects."""
    VALID_DEVICE_TYPE = ["WORKLOAD", "ENDPOINT"]
    VALID_OS_TYPE = ["CENTOS", "RHEL", "SLES", "UBUNTU", "WINDOWS"]
    VALID_SEVERITY = ["CRITICAL", "IMPORTANT", "MODERATE", "LOW"]
    VALID_SYNC_TYPE = ["MANUAL", "SCHEDULED"]
    VALID_SYNC_STATUS = ["NOT_STARTED", "MATCHED", "ERROR", "NOT_MATCHED",
                         "NOT_SUPPORTED", "CANCELLED", "IN_PROGRESS", "ACTIVE",
                         "COMPLETED"]
    VALID_DIRECTIONS = ["ASC", "DESC"]

    def __init__(self, doc_class, cb):
        """
        Initialize the VulnerabilityQuery.

        Args:
            doc_class (class): The model class that will be returned by this query.
            cb (BaseAPI): Reference to API object used to communicate with the server.
        """
        self._doc_class = doc_class
        self._cb = cb
        self._count_valid = False
        super(BaseQuery, self).__init__()
        self._query_builder = QueryBuilder()
        self._criteria = {}
        self._sortcriteria = {}
        self._total_results = 0
        self._vcenter_specific = None

    def set_vcenter(self, vcenter_id):
        """
        Restricts the vulnerabilities that this query is performed on to the specified vcenter id.

        Args:
            vcenter_id (string): vcenter id.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if vcenter_id:
            self._vcenter_specific = vcenter_id
        return self

    def set_device_type(self, device_type, operator):
        """
        Restricts the vulnerabilities that this query is performed on to the specified device type.

        Args:
            device_type (string): device type ("WORKLOAD", "ENDPOINT")
            operator (string):  logic operator to apply to property value.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if device_type not in VulnerabilityQuery.VALID_DEVICE_TYPE:
            raise ApiError("Invalid device type")
        self._update_criteria("device_type", device_type, operator)
        return self

    def set_highest_risk_score(self, highest_risk_score, operator):
        """
        Restricts the vulnerabilities that this query is performed on to the specified highest_risk_score.

        Args:
            highest_risk_score (double): highest_risk_score.
            operator (string):  logic operator to apply to property value.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if not highest_risk_score:
            raise ApiError("Invalid highest risk score")
        self._update_criteria("highest_risk_score", highest_risk_score, operator)
        return self

    def set_name(self, name, operator):
        """
        Restricts the vulnerabilities that this query is performed on to the specified name.

        Args:
            name (string): name.
            operator (string):  logic operator to apply to property value.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if not name:
            raise ApiError("Invalid name")
        self._update_criteria("name", name, operator)
        return self

    def set_os_arch(self, os_arch, operator):
        """
        Restricts the vulnerabilities that this query is performed on to the specified os_arch.

        Args:
            os_arch (string): os_arch.
            operator (string):  logic operator to apply to property value.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if not os_arch:
            raise ApiError("Invalid os architecture")
        self._update_criteria("os_arch", os_arch, operator)
        return self

    def set_os_name(self, os_name, operator):
        """
        Restricts the vulnerabilities that this query is performed on to the specified os_name.

        Args:
            os_name (string): os_name.
            operator (string):  logic operator to apply to property value.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if not os_name:
            raise ApiError("Invalid os name")
        self._update_criteria("os_name", os_name, operator)
        return self

    def set_os_type(self, os_type, operator):
        """
        Restricts the vulnerabilities that this query is performed on to the specified os type.

        Args:
            os_type (string): os type ("CENTOS", "RHEL", "SLES", "UBUNTU", "WINDOWS")
            operator (string):  logic operator to apply to property value.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if os_type not in VulnerabilityQuery.VALID_OS_TYPE:
            raise ApiError("Invalid os type")
        self._update_criteria("os_type", os_type, operator)
        return self

    def set_os_version(self, os_version, operator):
        """
        Restricts the vulnerabilities that this query is performed on to the specified os_version.

        Args:
            os_version (string): os_version.
            operator (string):  logic operator to apply to property value.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if not os_version:
            raise ApiError("Invalid os version")
        self._update_criteria("os_version", os_version, operator)
        return self

    def set_severity(self, severity, operator):
        """
        Restricts the vulnerabilities that this query is performed on to the specified severity.

        Args:
            severity (string): severity ("CRITICAL", "IMPORTANT", "MODERATE", "LOW")
            operator (string):  logic operator to apply to property value.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if severity not in VulnerabilityQuery.VALID_SEVERITY:
            raise ApiError("Invalid severity")
        self._update_criteria("severity", severity, operator)
        return self

    def set_sync_type(self, sync_type, operator):
        """
        Restricts the vulnerabilities that this query is performed on to the specified sync_type.

        Args:
            sync_type (string): sync_type ("MANUAL", "SCHEDULED")
            operator (string):  logic operator to apply to property value.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if sync_type not in VulnerabilityQuery.VALID_SYNC_TYPE:
            raise ApiError("Invalid sync type")
        self._update_criteria("sync_type", sync_type, operator)
        return self

    def set_sync_status(self, sync_status, operator):
        """
        Restricts the vulnerabilities that this query is performed on to the specified sync_status.

        Args:
            sync_status (string): sync_status ("NOT_STARTED", "MATCHED", "ERROR", "NOT_MATCHED", "NOT_SUPPORTED",
                                               "CANCELLED", "IN_PROGRESS", "ACTIVE", "COMPLETED")
            operator (string):  logic operator to apply to property value.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if sync_status not in VulnerabilityQuery.VALID_SYNC_STATUS:
            raise ApiError("Invalid sync status")
        self._update_criteria("sync_status", sync_status, operator)
        return self

    def set_vm_id(self, vm_id, operator):
        """
        Restricts the vulnerabilities that this query is performed on to the specified vm_id.

        Args:
            vm_id (string): vm_id.
            operator (string):  logic operator to apply to property value.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if not vm_id:
            raise ApiError("Invalid vm id")
        self._update_criteria("vm_id", vm_id, operator)
        return self

    def set_vuln_count(self, vuln_count, operator):
        """
        Restricts the vulnerabilities that this query is performed on to the specified vuln_count.

        Args:
            vuln_count (string): vuln_count.
            operator (string):  logic operator to apply to property value.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if not vuln_count:
            raise ApiError("Invalid vuln count")
        self._update_criteria("vuln_count", vuln_count, operator)
        return self

    def set_last_sync_ts(self, last_sync_ts, operator):
        """
        Restricts the vulnerabilities that this query is performed on to the specified last_sync_ts.

        Args:
            last_sync_ts (string): last_sync_ts.
            operator (string):  logic operator to apply to property value.

        Returns:
            VulnerabilityQuery: This instance.
        """
        if not last_sync_ts:
            raise ApiError("Invalid last_sync_ts")
        self._update_criteria("last_sync_ts", last_sync_ts, operator)
        return self

    """
        Including custom update_criteria, because the format is different:
        "criteria": {
            "property": {
                "value": "<string>",
                "operator": "<string>"
            }
        }
    """
    def _update_criteria(self, key, value, operator, overwrite=False):
        """
        Updates a list of criteria being collected for a query, by setting or appending items.

        Args:
            key (str): The property for the criteria item to be set.
            value (can be different types): the value for the criteria
            operator (str): any of the following types:
                - EQUALS, NOT_EQUALS, GREATER_THAN, LESS_THAN, IS_NULL,
                  IS_NOT_NULL, IS_TRUE, IS_FALSE, IN, NOT_IN, LIKE
            overwrite (bool): Overwrite the existing criteria for specified key

            The values are not lists, so if override is not allowed, disregard the change!
        """
        if self._criteria.get(key, None) is None or overwrite:
            self._criteria[key] = dict(value=value, operator=operator)

    def _build_request(self, from_row, max_rows, add_sort=True):
        """
        Creates the request body for an API call.

        Args:
            from_row (int): The row to start the query at.
            max_rows (int): The maximum number of rows to be returned.
            add_sort (bool): If True(default), the sort criteria will be added as part of the request.

        Returns:
            dict: The complete request body.
        """
        request = {"criteria": self._criteria, "query": self._query_builder._collapse(), "rows": 100}
        # Fetch 100 rows per page (instead of 10 by default) for better performance
        if from_row > 0:
            request["start"] = from_row
        if max_rows >= 0:
            request["rows"] = max_rows
        if add_sort and self._sortcriteria != {}:
            request["sort"] = [self._sortcriteria]
        return request

    def _build_url(self, tail_end):
        """
        Creates the URL to be used for an API call.

        Args:
            tail_end (str): String to be appended to the end of the generated URL.

        Returns:
            str: The complete URL.
        """
        if self._vcenter_specific:
            additional = "/vcenters/{}/devices/vulnerabilities".format(self._vcenter_specific)
        else:
            additional = "/devices/vulnerabilities"
        additional += self._doc_class.url_additional
        url = self._doc_class.urlobject.format(self._cb.credentials.org_key) + additional + tail_end
        return url

    def _count(self):
        """
        Returns the number of results from the run of this query.

        Returns:
            int: The number of results from the run of this query.
        """
        if self._count_valid:
            return self._total_results
        url = self._build_url("/_search")
        request = self._build_request(0, -1)
        resp = self._cb.post_object(url, body=request)
        result = resp.json()

        self._total_results = result["num_found"]
        self._count_valid = True

        return self._total_results

    def _perform_query(self, from_row=0, max_rows=-1):
        """
        Performs the query and returns the results of the query in an iterable fashion.

        Args:
            from_row (int): The row to start the query at (default 0).
            max_rows (int): The maximum number of rows to be returned (default -1, meaning "all").

        Returns:
            Iterable: The iterated query.
        """
        url = self._build_url("/_search")
        current = from_row
        numrows = 0
        still_querying = True
        while still_querying:
            request = self._build_request(current, max_rows)
            resp = self._cb.post_object(url, body=request)
            result = resp.json()

            self._total_results = result["num_found"]
            self._count_valid = True

            results = result.get("results", [])
            for item in results:
                yield self._doc_class(self._cb, initial_data=item)
                current += 1
                numrows += 1

                if max_rows > 0 and numrows == max_rows:
                    still_querying = False
                    break

            if current >= self._total_results:
                break

    def _run_async_query(self, context):
        """
        Executed in the background to run an asynchronous query.

        Args:
            context (object): Not used, always None.

        Returns:
            list: Result of the async query, which is then returned by the future.
        """
        url = self._build_url("/_search")
        request = self._build_request(0, -1)
        resp = self._cb.post_object(url, body=request)
        result = resp.json()
        self._total_results = result["num_found"]
        self._count_valid = True
        results = result.get("results", [])
        return [self._doc_class(self._cb, initial_data=item) for item in results]

    def sort_by(self, key, direction="ASC"):
        """
        Sets the sorting behavior on a query's results.

        Example:
            >>> cb.select(Vulnerabiltiy).sort_by("status")

        Args:
            key (str): The key in the schema to sort by.
            direction (str): The sort order, either "ASC" or "DESC".

        Returns:
            VulnerabilityQuery: This instance.

        Raises:
            ApiError: If an invalid direction value is passed.
        """
        if direction not in VulnerabilityQuery.VALID_DIRECTIONS:
            raise ApiError("invalid sort direction specified")
        self._sortcriteria = {"field": key, "order": direction}
        return self


class AffectedAssetQuery(VulnerabilityQuery):
    """Query Class for the Vulnerability"""

    def __init__(self, vulnerability, cb):
        """
        Initialize the AffectedAssetQuery.

        Args:
            vulnerability (class): The vulnerability that will be returned by this query.
            cb (BaseAPI): Reference to API object used to communicate with the server.
        """
        # to avoid circular imports
        from cbc_sdk.platform import Device
        self.vulnerability = vulnerability
        self._doc_class = Device
        self._cb = cb
        self._count_valid = False
        super(BaseQuery, self).__init__()
        self._query_builder = QueryBuilder()
        self._criteria = {}
        self._sortcriteria = {}
        self._total_results = 0
        self._vcenter_specific = None

    def _build_url(self, tail_end):
        """
        Creates the URL to be used for an API call.

        Args:
            tail_end (str): String to be appended to the end of the generated URL.

        Returns:
            str: The complete URL.
        """
        if self._vcenter_specific:
            additional = "/vcenters/{}/vulnerabilities/{}/devices".format(self._vcenter_specific,
                                                                          self.vulnerability._model_unique_id)
        else:
            additional = "/vulnerabilities/{}/devices".format(self.vulnerability._model_unique_id)
        additional += self.vulnerability.url_additional
        url = (self.vulnerability.urlobject + '/orgs/{}').format(self._cb.credentials.org_key) + additional + tail_end
        return url

    def _build_request(self, from_row, max_rows, add_sort=True):
        """
        Creates the request body for an API call.

        Args:
            from_row (int): The row to start the query at.
            max_rows (int): The maximum number of rows to be returned.
            add_sort (bool): If True(default), the sort criteria will be added as part of the request.

        Returns:
            dict: The complete request body.
        """
        request = {
            "os_product_id": self.vulnerability._os_product_id,
            "criteria": self._criteria,
            "query": self._query_builder._collapse(),
            "rows": 100}
        # Fetch 100 rows per page (instead of 10 by default) for better performance
        if from_row > 0:
            request["start"] = from_row
        if max_rows >= 0:
            request["rows"] = max_rows
        if add_sort and self._sortcriteria != {}:
            request["sort"] = [self._sortcriteria]
        return request

    def _perform_query(self, from_row=0, max_rows=-1):
        """
        Performs the query and returns the results of the query in an iterable fashion.

        Args:
            from_row (int): The row to start the query at (default 0).
            max_rows (int): The maximum number of rows to be returned (default -1, meaning "all").

        Returns:
            Iterable: The iterated query.
        """
        url = self._build_url("")
        current = from_row
        numrows = 0
        still_querying = True
        while still_querying:
            request = self._build_request(current, max_rows)
            resp = self._cb.post_object(url, body=request)
            result = resp.json()

            self._total_results = result["num_found"]
            self._count_valid = True

            results = result.get("results", [])
            # avoid circular import error
            from cbc_sdk.platform import Device
            device_ids = [asset["device_id"] for asset in results]
            # this would load all of the information for the devices
            # check whether it is empty list, because otherwise all devices are returned
            if device_ids:
                devices = self._cb.select(Device).set_device_ids(device_ids)
                for device in devices:
                    yield device
                    current += 1
                    numrows += 1

                    if max_rows > 0 and numrows == max_rows:
                        still_querying = False
                        break

            if current >= self._total_results:
                break

    def _run_async_query(self, context):
        """
        Executed in the background to run an asynchronous query.

        Args:
            context (object): Not used, always None.

        Returns:
            list: Result of the async query, which is then returned by the future.
        """
        url = self._build_url("")
        request = self._build_request(0, -1)
        resp = self._cb.post_object(url, body=request)
        result = resp.json()
        self._total_results = result["num_found"]
        self._count_valid = True
        results = result.get("results", [])
        # to avoid circular import error
        from cbc_sdk.platform import Device
        # returning all of the information for the device
        # check whether there are device ids, because otherwise all of the devices would be returned
        device_ids = [item['device_id'] for item in results]
        if device_ids:
            return self._cb.select(Device).set_device_ids(device_ids)
        else:
            return []

    def _count(self):
        """
        Returns the number of results from the run of this query.

        Returns:
            int: The number of results from the run of this query.
        """
        if self._count_valid:
            return self._total_results
        url = self._build_url("")
        request = self._build_request(0, -1)
        resp = self._cb.post_object(url, body=request)
        result = resp.json()

        self._total_results = result["num_found"]
        self._count_valid = True

        return self._total_results


class DeviceVulnerabilityQuery(VulnerabilityQuery):
    """Query Class for the DeviceVulnerability"""

    def __init__(self, device, cb):
        """
        Initialize the DeviceVulnerabilityQuery.

        Args:
            device (class): The model class (Device) that will be returned by this query.
            cb (BaseAPI): Reference to API object used to communicate with the server.
        """
        self.device = device
        self._doc_class = DeviceVulnerability
        self._cb = cb
        self._count_valid = False
        super(BaseQuery, self).__init__()
        self._query_builder = QueryBuilder()
        self._criteria = {}
        self._sortcriteria = {}
        self._total_results = 0

    def _build_url(self, tail_end):
        """
        Creates the URL to be used for an API call.

        Args:
            tail_end (str): String to be appended to the end of the generated URL.

        Returns:
            str: The complete URL.
        """
        urlobject = '/vulnerability/assessment/api/v1/orgs/{}'.format(self._cb.credentials.org_key)
        if self.device._vcenter_specific:
            urlobject += "/vcenters/{}".format(self.device._vcenter_specific)
        urlobject += '/devices/{}/vulnerabilities'.format(self.device._model_unique_id)
        url = urlobject + tail_end
        return url

    def _build_request(self, from_row, max_rows, add_sort=True):
        """
        Creates the request body for an API call.

        Args:
            from_row (int): The row to start the query at.
            max_rows (int): The maximum number of rows to be returned.
            add_sort (bool): If True(default), the sort criteria will be added as part of the request.

        Returns:
            dict: The complete request body.
        """
        request = {
            "criteria": self._criteria,
            "query": self._query_builder._collapse(),
            "rows": 100}
        # Fetch 100 rows per page (instead of 10 by default) for better performance
        if from_row > 0:
            request["start"] = from_row
        if max_rows >= 0:
            request["rows"] = max_rows
        if add_sort and self._sortcriteria != {}:
            request["sort"] = [self._sortcriteria]
        return request

    def _perform_query(self, from_row=0, max_rows=-1):
        """
        Performs the query and returns the results of the query in an iterable fashion.

        Args:
            from_row (int): The row to start the query at (default 0).
            max_rows (int): The maximum number of rows to be returned (default -1, meaning "all").

        Returns:
            Iterable: The iterated query.
        """
        url = self._build_url("/_search")
        current = from_row
        numrows = 0
        still_querying = True
        while still_querying:
            request = self._build_request(current, max_rows)
            resp = self._cb.post_object(url, body=request)
            result = resp.json()

            self._total_results = result["num_found"]
            self._count_valid = True

            results = result.get("results", [])
            for item in results:
                yield self._doc_class(self._cb, initial_data=item)
                current += 1
                numrows += 1

                if max_rows > 0 and numrows == max_rows:
                    still_querying = False
                    break

            if current >= self._total_results:
                break

    def _run_async_query(self, context):
        """
        Executed in the background to run an asynchronous query.

        Args:
            context (object): Not used, always None.

        Returns:
            list: Result of the async query, which is then returned by the future.
        """
        url = self._build_url("/_search")
        request = self._build_request(0, -1)
        resp = self._cb.post_object(url, body=request)
        result = resp.json()
        self._total_results = result["num_found"]
        self._count_valid = True
        results = result.get("results", [])
        return [self._doc_class(self._cb, initial_data=item) for item in results]

    def _count(self):
        """
        Returns the number of results from the run of this query.

        Returns:
            int: The number of results from the run of this query.
        """
        if self._count_valid:
            return self._total_results
        url = self._build_url("/_search")
        request = self._build_request(0, -1)
        resp = self._cb.post_object(url, body=request)
        result = resp.json()

        self._total_results = result["num_found"]
        self._count_valid = True

        return self._total_results
