<commit_before>import StringIO
import csv
import json
import logging
import shutil
from collections import namedtuple
from stat import S_IRUSR

import pytest
from retrying import retry
from subprocess import Popen

import mazerunner
import os

from mazerunner.api_client import DeploymentGroupCollection, \
    BreadcrumbCollection, ServiceCollection, DecoyCollection, Service, \
    AlertPolicy, CIDRMappingCollection, BackgroundTaskCollection, \
    EndpointCollection, Decoy, Breadcrumb, DeploymentGroup, Endpoint, CIDRMapping, BackgroundTask
from mazerunner.exceptions import ValidationError, ServerError, BadParamError, \
    InvalidInstallMethodError
from utils import TimeoutException, wait_until

ENDPOINT_IP_PARAM = 'endpoint_ip'
ENDPOINT_USERNAME_PARAM = 'endpoint_username'
ENDPOINT_PASSWORD_PARAM = '<PASSWORD>'

CODE_EXECUTION_ALERT_TYPE = 'code'
FORENSIC_DATA_ALERT_TYPE = 'forensic_puller'

MAZERUNNER_IP_ADDRESS_PARAM = 'ip_address'
API_ID_PARAM = 'id'
API_SECRET_PARAM = 'secret'
MAZERUNNER_CERTIFICATE_PATH_PARAM = 'mazerunner_certificate_path'

ENTITIES_CONFIGURATION = {
    Decoy: [],
    Service: [],
    Breadcrumb: [],
    DeploymentGroup: [1],
    Endpoint: [],
    CIDRMapping: [],
    BackgroundTask: []
}


TEST_DEPLOYMENTS_FILE_PATH = os.path.join(os.path.dirname(__file__), 'test_deployments/dep.zip')
TEST_DEPLOYMENTS_FOLDER_PATH = os.path.dirname(TEST_DEPLOYMENTS_FILE_PATH)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("APITest")
logger.setLevel(logging.INFO)


class AlertNotFoundError(RuntimeError):
    pass


def _clear_deployment_path():
    if os.path.exists(TEST_DEPLOYMENTS_FOLDER_PATH):
        shutil.rmtree(TEST_DEPLOYMENTS_FOLDER_PATH)

    os.makedirs(TEST_DEPLOYMENTS_FOLDER_PATH)


class MachineStatus(object):
    NOT_SEEN = "not_seen"
    ACTIVE = "active"
    INACTIVE = "inactive"


# noinspection PyMethodMayBeStatic,PyAttributeOutsideInit
class APITest(object):

    runslow = pytest.mark.skipif(not pytest.config.getoption('--runslow'),
                                 reason='--runslow not activated')
    lab_dependent = pytest.mark.skipif(not pytest.config.getoption('--lab_dependent'),
                                       reason='--lab_dependent not activated')

    def _assert_clean_system(self):
        for entity_collection in self.disposable_entities:
            existing_ids = {entity.id for entity in entity_collection}
            expected_ids = set(ENTITIES_CONFIGURATION[entity_collection.MODEL_CLASS])

            assert existing_ids == expected_ids, 'System must be clean before running this test. ' \
                                                 'Use the --initial_clean flag to do this ' \
                                                 'automatically'

    def _configure_entities_groups(self):
        self.decoys = self.client.decoys
        self.services = self.client.services
        self.breadcrumbs = self.client.breadcrumbs
        self.deployment_groups = self.client.deployment_groups
        self.alerts = self.client.alerts
        self.alert_policies = self.client.alert_policies
        self.cidr_mappings = self.client.cidr_mappings
        self.endpoints = self.client.endpoints
        self.background_tasks = self.client.background_tasks

        self.disposable_entities = [
            self.decoys,
            self.services,
            self.breadcrumbs,
            self.deployment_groups,
            self.endpoints,
            self.cidr_mappings,
            self.background_tasks
        ]

    def setup_method(self, method):
        logger.debug("setup_method called")

        with open(pytest.config.option.json_credentials, 'rb') as file_reader:
            json_dict = json.load(file_reader)

        self.lab_endpoint_ip = json_dict.get(ENDPOINT_IP_PARAM)
        self.lab_endpoint_user = json_dict.get(ENDPOINT_USERNAME_PARAM)
        self.lab_endpoint_password = json_dict.get(ENDPOINT_PASSWORD_PARAM)

        self.mazerunner_ip_address = json_dict[MAZERUNNER_IP_ADDRESS_PARAM]
        self.api_key = json_dict[API_ID_PARAM]
        self.api_secret = json_dict[API_SECRET_PARAM]
        self.mazerunner_certificate_path = json_dict[MAZERUNNER_CERTIFICATE_PATH_PARAM]

        self.client = mazerunner.connect(
            ip_address=self.mazerunner_ip_address,
            api_key=self.api_key,
            api_secret=self.api_secret,
            certificate=self.mazerunner_certificate_path)

        self._configure_entities_groups()

        if pytest.config.option.initial_clean:
            self._destroy_new_entities()

        self._assert_clean_system()

        self.file_paths_for_cleanup = []

        _clear_deployment_path()

    def _destroy_new_entities(self):
        for entity_collection in self.disposable_entities:
            for entity in list(entity_collection):
                initial_ids = ENTITIES_CONFIGURATION[entity_collection.MODEL_CLASS]
                if entity.id not in initial_ids:
                    wait_until(entity.delete, exc_list=[ServerError, ValidationError],
                               check_return_value=False)

        wait_until(self._assert_clean_system, exc_list=[AssertionError], check_return_value=False)

    def teardown_method(self, method):
        logger.debug("teardown_method called")

        self._destroy_new_entities()

        # Clean files:
        for file_path in self.file_paths_for_cleanup:
            if os.path.exists(file_path):
                os.remove(file_path)
        _clear_deployment_path()

    def valid_decoy_status(self, decoy, wanted_statuses):
        logger.debug("valid_decoy_status called")
        decoy.load()
        return decoy.machine_status in wanted_statuses

    def wait_for_decoy_status(self, decoy, wanted_statuses, timeout):
        logger.info("wait_for_decoy_status called")
        logger.info("waiting up to %d seconds", timeout)
        try:
            wait_until(
                self.valid_decoy_status,
                decoy=decoy,
                wanted_statuses=wanted_statuses,
                check_return_value=True,<commit_msg>Fix how background tasks are cleared in tests<commit_after>import StringIO
import csv
import json
import logging
import shutil
from collections import namedtuple
from stat import S_IRUSR

import pytest
from retrying import retry
from subprocess import Popen

import mazerunner
import os

from mazerunner.api_client import DeploymentGroupCollection, \
    BreadcrumbCollection, ServiceCollection, DecoyCollection, Service, \
    AlertPolicy, CIDRMappingCollection, BackgroundTaskCollection, \
    EndpointCollection, Decoy, Breadcrumb, DeploymentGroup, Endpoint, CIDRMapping, BackgroundTask
from mazerunner.exceptions import ValidationError, ServerError, BadParamError, \
    InvalidInstallMethodError
from utils import TimeoutException, wait_until

CLEAR_SYSTEM_ERROR_MESSAGE = 'System must be clean before running this test. Use the '\
                             '--initial_clean flag to do this automatically'
ENDPOINT_IP_PARAM = 'endpoint_ip'
ENDPOINT_USERNAME_PARAM = 'endpoint_username'
ENDPOINT_PASSWORD_PARAM = 'endpoint_password'

CODE_EXECUTION_ALERT_TYPE = 'code'
FORENSIC_DATA_ALERT_TYPE = 'forensic_puller'

MAZERUNNER_IP_ADDRESS_PARAM = 'ip_address'
API_ID_PARAM = 'id'
API_SECRET_PARAM = 'secret'
MAZERUNNER_CERTIFICATE_PATH_PARAM = 'mazerunner_certificate_path'

ENTITIES_CONFIGURATION = {
    Decoy: [],
    Service: [],
    Breadcrumb: [],
    DeploymentGroup: [1],
    Endpoint: [],
    CIDRMapping: [],
    BackgroundTask: []
}


TEST_DEPLOYMENTS_FILE_PATH = os.path.join(os.path.dirname(__file__), 'test_deployments/dep.zip')
TEST_DEPLOYMENTS_FOLDER_PATH = os.path.dirname(TEST_DEPLOYMENTS_FILE_PATH)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("APITest")
logger.setLevel(logging.INFO)


class AlertNotFoundError(RuntimeError):
    pass


def _clear_deployment_path():
    if os.path.exists(TEST_DEPLOYMENTS_FOLDER_PATH):
        shutil.rmtree(TEST_DEPLOYMENTS_FOLDER_PATH)

    os.makedirs(TEST_DEPLOYMENTS_FOLDER_PATH)


class MachineStatus(object):
    NOT_SEEN = "not_seen"
    ACTIVE = "active"
    INACTIVE = "inactive"


# noinspection PyMethodMayBeStatic,PyAttributeOutsideInit
class APITest(object):

    runslow = pytest.mark.skipif(not pytest.config.getoption('--runslow'),
                                 reason='--runslow not activated')
    lab_dependent = pytest.mark.skipif(not pytest.config.getoption('--lab_dependent'),
                                       reason='--lab_dependent not activated')

    def _assert_clean_system(self):
        for entity_collection in self.disposable_entities:
            existing_ids = {entity.id for entity in entity_collection}
            expected_ids = set(ENTITIES_CONFIGURATION[entity_collection.MODEL_CLASS])

            assert existing_ids == expected_ids, CLEAR_SYSTEM_ERROR_MESSAGE

        assert len(self.background_tasks) == 0, CLEAR_SYSTEM_ERROR_MESSAGE

    def _configure_entities_groups(self):
        self.decoys = self.client.decoys
        self.services = self.client.services
        self.breadcrumbs = self.client.breadcrumbs
        self.deployment_groups = self.client.deployment_groups
        self.alerts = self.client.alerts
        self.alert_policies = self.client.alert_policies
        self.cidr_mappings = self.client.cidr_mappings
        self.endpoints = self.client.endpoints
        self.background_tasks = self.client.background_tasks

        self.disposable_entities = [
            self.decoys,
            self.services,
            self.breadcrumbs,
            self.deployment_groups,
            self.endpoints,
            self.cidr_mappings
        ]

    def setup_method(self, method):
        logger.debug("setup_method called")

        with open(pytest.config.option.json_credentials, 'rb') as file_reader:
            json_dict = json.load(file_reader)

        self.lab_endpoint_ip = json_dict.get(ENDPOINT_IP_PARAM)
        self.lab_endpoint_user = json_dict.get(ENDPOINT_USERNAME_PARAM)
        self.lab_endpoint_password = json_dict.get(ENDPOINT_PASSWORD_PARAM)

        self.mazerunner_ip_address = json_dict[MAZERUNNER_IP_ADDRESS_PARAM]
        self.api_key = json_dict[API_ID_PARAM]
        self.api_secret = json_dict[API_SECRET_PARAM]
        self.mazerunner_certificate_path = json_dict[MAZERUNNER_CERTIFICATE_PATH_PARAM]

        self.client = mazerunner.connect(
            ip_address=self.mazerunner_ip_address,
            api_key=self.api_key,
            api_secret=self.api_secret,
            certificate=self.mazerunner_certificate_path)

        self._configure_entities_groups()

        if pytest.config.option.initial_clean:
            self._destroy_new_entities()

        self._assert_clean_system()

        self.file_paths_for_cleanup = []

        _clear_deployment_path()

    def _destroy_new_entities(self):
        for entity_collection in self.disposable_entities:
            for entity in list(entity_collection):
                initial_ids = ENTITIES_CONFIGURATION[entity_collection.MODEL_CLASS]
                if entity.id not in initial_ids:
                    wait_until(entity.delete, exc_list=[ServerError, ValidationError],
                               check_return_value=False)

        self.background_tasks.acknowledge_all_complete()

        wait_until(self._assert_clean_system, exc_list=[AssertionError], check_return_value=False)

    def teardown_method(self, method):
        logger.debug("teardown_method called")

        self._destroy_new_entities()

        # Clean files:
        for file_path in self.file_paths_for_cleanup:
            if os.path.exists(file_path):
                os.remove(file_path)
        _clear_deployment_path()

    def valid_decoy_status(self, decoy, wanted_statuses):
        logger.debug("valid_decoy_status called")
        decoy.load()
        return decoy.machine_status in wanted_statuses

    def wait_for_decoy_status(self, decoy, wanted_statuses, timeout):
        logger.info("wait_for_decoy_status called")
        logger.info("waiting up to %d seconds", timeout)
        try:
            wait_until(
                self.valid_decoy_status,
                decoy=decoy,
                wanted_statuses=wanted_statuses,
                check_return_value=True,
{'fraction_non_alphanumeric': 0.07601721757492082, 'fraction_numerical': 0.0002436449281247462, 'mean_word_length': 3.106035345115038, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1498935', 'n_tokens_mistral': 3806, 'n_tokens_neox': 3705, 'n_words': 614}