def rescan(self, submission: Submission, results: Dict[str, Result],
               file_infos: Dict[str, FileInfo], file_tree, errors: List[str],
               rescan_services: List[str]):
        """
        Rescan a submission started on another system.
        """
        # Reset submission processing data
        submission['times'].pop('completed')
        submission['state'] = 'submitted'

        # Set the list of service to rescan
        submission['params']['services']['rescan'] = rescan_services

        # Create the submission object
        submission_obj = Submission(submission)

        if len(submission_obj.files) == 0:
            raise SubmissionException("No files found to submit.")

        for f in submission_obj.files:
            if not self.datastore.file.exists(f.sha256):
                raise SubmissionException(
                    f"File {f.sha256} does not exist, cannot continue submission."
                )

        # Set the new expiry
        if submission_obj.params.ttl:
            submission_obj.expiry_ts = epoch_to_iso(now() +
                                                    submission_obj.params.ttl *
                                                    24 * 60 * 60)

        # Clearing runtime_excluded on initial submit or resubmit
        submission_obj.params.services.runtime_excluded = []

        # Save the submission
        self.datastore.submission.save(submission_obj.sid, submission_obj)

        # Dispatch the submission
        self.log.debug("Submission complete. Dispatching: %s",
                       submission_obj.sid)
        self.dispatcher.dispatch_bundle(submission_obj, results, file_infos,
                                        file_tree, errors)

        return submission
    def dispatch_submission(self, submission: Submission, completed_queue: str = None):
        """Insert a submission into the dispatching system.

        Note:
            You probably actually want to use the SubmissionTool

        Prerequsits:
            - submission should already be saved in the datastore
            - files should already be in the datastore and filestore
        """
        self.submission_queue.push(dict(
            submission=submission.as_primitives(),
            completed_queue=completed_queue,
        ))
    def dispatch_bundle(self, submission: Submission, results: Dict[str, Result],
                        file_infos: Dict[str, File], file_tree, errors: Dict[str, Error], completed_queue: str = None):
        """Insert a bundle into the dispatching system and continue scanning of its files

        Prerequisites:
            - Submission, results, file_infos and errors should already be saved in the datastore
            - Files should already be in the filestore
        """
        self.submission_queue.push(dict(
            submission=submission.as_primitives(),
            results=results,
            file_infos=file_infos,
            file_tree=file_tree,
            errors=errors,
            completed_queue=completed_queue,
        ))
    def submit(self,
               submission_obj: SubmissionObject,
               local_files: List = None,
               completed_queue=None):
        """Submit several files in a single submission.

        After this method runs, there should be no local copies of the file left.
        """
        if local_files is None:
            local_files = []

        if len(submission_obj.files) == 0 and len(local_files) == 0:
            raise SubmissionException("No files found to submit...")

        if submission_obj.params.ttl:
            expiry = epoch_to_iso(submission_obj.time.timestamp() +
                                  submission_obj.params.ttl * 24 * 60 * 60)
        else:
            expiry = None
        max_size = self.config.submission.max_file_size

        for local_file in local_files:
            # Upload/download, extract, analyze files
            original_classification = str(submission_obj.params.classification)
            file_hash, size, new_metadata = self._ready_file(
                local_file, expiry, original_classification)
            new_name = new_metadata.pop('name',
                                        safe_str(os.path.basename(local_file)))
            meta_classification = new_metadata.pop('classification',
                                                   original_classification)
            if meta_classification != original_classification:
                try:
                    submission_obj.params.classification = Classification.max_classification(
                        meta_classification, original_classification)
                except InvalidClassification as ic:
                    raise SubmissionException(
                        "The classification found inside the cart file cannot be merged with "
                        f"the classification the file was submitted as: {str(ic)}"
                    )

            submission_obj.metadata.update(**flatten(new_metadata))

            # Check that after we have resolved exactly what to pass on, that it
            # remains a valid target for scanning
            if size > max_size and not submission_obj.params.ignore_size:
                msg = "File too large (%d > %d). Submission failed" % (
                    size, max_size)
                raise SubmissionException(msg)
            elif size == 0:
                msg = "File empty. Submission failed"
                raise SubmissionException(msg)

            submission_obj.files.append(
                File({
                    'name': new_name,
                    'size': size,
                    'sha256': file_hash,
                }))

        # Clearing runtime_excluded on initial submit or resubmit
        submission_obj.params.services.runtime_excluded = []

        # We should now have all the information we need to construct a submission object
        sub = Submission(
            dict(
                archive_ts=now_as_iso(
                    self.config.datastore.ilm.days_until_archive * 24 * 60 *
                    60),
                classification=submission_obj.params.classification,
                error_count=0,
                errors=[],
                expiry_ts=expiry,
                file_count=len(submission_obj.files),
                files=submission_obj.files,
                max_score=0,
                metadata=submission_obj.metadata,
                params=submission_obj.params,
                results=[],
                sid=submission_obj.sid,
                state='submitted',
                scan_key=submission_obj.scan_key,
            ))

        if self.config.ui.allow_malicious_hinting and submission_obj.params.malicious:
            sub.verdict = {"malicious": [submission_obj.params.submitter]}

        self.datastore.submission.save(sub.sid, sub)

        self.log.debug("Submission complete. Dispatching: %s", sub.sid)
        self.dispatcher.dispatch_submission(sub,
                                            completed_queue=completed_queue)

        return sub
def test_update_single_alert(config, datastore):
    persistent_redis = get_client(
        host=config.core.redis.persistent.host,
        port=config.core.redis.persistent.port,
        private=False,
    )
    alerter = Alerter()
    # Swap our alerter onto a private queue so our test doesn't get intercepted
    alerter.alert_queue = alert_queue = NamedQueue(uuid.uuid4().hex,
                                                   persistent_redis)

    # Get a random submission
    submission = random.choice(all_submissions)
    all_submissions.remove(submission)

    # Generate a task for the submission
    ingest_msg = random_model_obj(IngestTask)
    ingest_msg.submission.sid = submission.sid
    ingest_msg.submission.metadata = submission.metadata
    ingest_msg.submission.params = submission.params
    ingest_msg.submission.files = submission.files

    alert_queue.push(ingest_msg.as_primitives())
    alert_type = alerter.run_once()
    assert alert_type == 'create'
    datastore.alert.commit()

    original_alert = datastore.alert.get(
        datastore.alert.search(f"sid:{submission.sid}", fl="id",
                               as_obj=False)['items'][0]['id'])
    assert original_alert is not None

    # Generate a children task
    child_submission = Submission(submission.as_primitives())
    child_submission.sid = get_random_id()
    child_submission.params.psid = submission.sid

    # Alter the result of one of the services
    r = None
    while r is None:
        r = datastore.result.get(random.choice(child_submission.results))

    for s in r.result.sections:
        old_tags = s.tags.as_primitives(strip_null=True)
        s.tags = Tagging(recursive_extend(old_tags, get_random_tags()))

    datastore.result.save(r.build_key(), r)
    datastore.result.commit()

    datastore.submission.save(child_submission.sid, child_submission)
    datastore.submission.commit()

    child_ingest_msg = random_model_obj(IngestTask)
    child_ingest_msg.submission.sid = child_submission.sid
    child_ingest_msg.submission.metadata = child_submission.metadata
    child_ingest_msg.submission.params = child_submission.params
    child_ingest_msg.submission.files = child_submission.files
    child_ingest_msg.submission.time = ingest_msg.submission.time

    alert_queue.push(child_ingest_msg.as_primitives())
    alert_type = alerter.run_once()
    assert alert_type == 'update'
    datastore.alert.commit()

    updated_alert = datastore.alert.get(
        datastore.alert.search(f"sid:{child_submission.sid}",
                               fl="id",
                               as_obj=False)['items'][0]['id'])
    assert updated_alert is not None

    assert updated_alert != original_alert
    def try_run(self, volatile=False):
        ingester = self.ingester
        cpu_mark = time.process_time()
        time_mark = time.time()

        # Move from ingest to unique and waiting queues.
        # While there are entries in the ingest queue we consume chunk_size
        # entries at a time and move unique entries to uniqueq / queued and
        # duplicates to their own queues / waiting.
        while self.running:
            self.heartbeat()
            while True:
                result = ingester.complete_queue.pop(blocking=False)
                if not result:
                    break
                # Start of ingest message
                if self.apm_client:
                    self.apm_client.begin_transaction('ingest_msg')

                sub = Submission(result)
                ingester.completed(sub)

                # End of ingest message (success)
                if self.apm_client:
                    elasticapm.tag(sid=sub.sid)
                    self.apm_client.end_transaction('ingest_complete', 'success')

            ingester.counter.increment_execution_time('cpu_seconds', time.process_time() - cpu_mark)
            ingester.counter.increment_execution_time('busy_seconds', time.time() - time_mark)

            message = ingester.ingest_queue.pop(timeout=1)

            cpu_mark = time.process_time()
            time_mark = time.time()

            if not message:
                continue

            # Start of ingest message
            if self.apm_client:
                self.apm_client.begin_transaction('ingest_msg')

            try:
                sub = SubmissionInput(message)
                # Write all input to the traffic queue
                ingester.traffic_queue.publish(SubmissionMessage({
                    'msg': sub,
                    'msg_type': 'SubmissionIngested',
                    'sender': 'ingester',
                }).as_primitives())

                task = IngestTask(dict(
                    submission=sub,
                    ingest_id=sub.sid,
                ))
                task.submission.sid = None  # Reset to new random uuid

            except (ValueError, TypeError) as error:
                self.log.exception(f"Dropped ingest submission {message} because {str(error)}")

                # End of ingest message (value_error)
                if self.apm_client:
                    self.apm_client.end_transaction('ingest_input', 'value_error')

                if volatile:
                    raise
                continue

            if any(len(file.sha256) != 64 for file in task.submission.files):
                self.log.error(f"Invalid sha256: {[file.sha256 for file in task.submission.files]}")

                # End of ingest message (invalid_hash)
                if self.apm_client:
                    self.apm_client.end_transaction('ingest_input', 'invalid_hash')

                continue

            for file in task.submission.files:
                file.sha256 = file.sha256.lower()

            ingester.ingest(task)

            # End of ingest message (success)
            if self.apm_client:
                self.apm_client.end_transaction('ingest_input', 'success')
Esempio n. 7
0
    def submit(self, submission_obj: SubmissionObject, local_files: List = None, cleanup=True, completed_queue=None):
        """Submit several files in a single submission.

        After this method runs, there should be no local copies of the file left.
        """
        if local_files is None:
            local_files = []

        try:
            expiry = now_as_iso(submission_obj.params.ttl * 24 * 60 * 60) if submission_obj.params.ttl else None
            max_size = self.config.submission.max_file_size

            if len(submission_obj.files) == 0:
                if len(local_files) == 0:
                    raise SubmissionException("No files found to submit...")

                for local_file in local_files:
                    # Upload/download, extract, analyze files
                    file_hash, size, new_metadata = self._ready_file(local_file, expiry,
                                                                     str(submission_obj.params.classification),
                                                                     cleanup, upload=True)
                    new_name = new_metadata.pop('name', safe_str(os.path.basename(local_file)))
                    submission_obj.params.classification = new_metadata.pop('classification',
                                                                            submission_obj.params.classification)
                    submission_obj.metadata.update(**flatten(new_metadata))

                    # Check that after we have resolved exactly what to pass on, that it
                    # remains a valid target for scanning
                    if size > max_size and not submission_obj.params.ignore_size:
                        msg = "File too large (%d > %d). Submission failed" % (size, max_size)
                        raise SubmissionException(msg)
                    elif size == 0:
                        msg = "File empty. Submission failed"
                        raise SubmissionException(msg)

                    submission_obj.files.append(File({
                        'name': new_name,
                        'size': size,
                        'sha256': file_hash,
                    }))
            else:
                for f in submission_obj.files:
                    temporary_path = None
                    try:
                        fd, temporary_path = tempfile.mkstemp(prefix="submission.submit")
                        os.close(fd)  # We don't need the file descriptor open
                        self.filestore.download(f.sha256, temporary_path)
                        file_hash, size, new_metadata = self._ready_file(temporary_path, expiry,
                                                                         str(submission_obj.params.classification),
                                                                         cleanup, sha256=f.sha256)

                        new_name = new_metadata.pop('name', f.name)
                        submission_obj.params.classification = new_metadata.pop('classification',
                                                                                submission_obj.params.classification)
                        submission_obj.metadata.update(**flatten(new_metadata))

                        # Check that after we have resolved exactly what to pass on, that it
                        # remains a valid target for scanning
                        if size > max_size and not submission_obj.params.ignore_size:
                            msg = "File too large (%d > %d). Submission failed" % (size, max_size)
                            raise SubmissionException(msg)
                        elif size == 0:
                            msg = "File empty. Submission failed"
                            raise SubmissionException(msg)

                        if f.size is None:
                            f.size = size

                        f.name = new_name
                        f.sha256 = file_hash

                    finally:
                        if temporary_path:
                            if os.path.exists(temporary_path):
                                os.unlink(temporary_path)

            # Initialize the temporary data from the submission parameter
            if submission_obj.params.initial_data:
                try:
                    temp_hash_name = get_temporary_submission_data_name(submission_obj.sid,
                                                                        submission_obj.files[0].sha256)
                    temporary_submission_data = ExpiringHash(temp_hash_name, host=self.redis)
                    temporary_submission_data.multi_set(json.loads(submission_obj.params.initial_data))
                except ValueError as err:
                    self.log.warning(f"[{submission_obj.sid}] could not process initialization data: {err}")

            # Clearing runtime_excluded on initial submit or resubmit
            submission_obj.params.services.runtime_excluded = []

            # We should now have all the information we need to construct a submission object
            sub = Submission(dict(
                archive_ts=now_as_iso(self.config.datastore.ilm.days_until_archive * 24 * 60 * 60),
                classification=submission_obj.params.classification,
                error_count=0,
                errors=[],
                expiry_ts=expiry,
                file_count=len(submission_obj.files),
                files=submission_obj.files,
                max_score=0,
                metadata=submission_obj.metadata,
                params=submission_obj.params,
                results=[],
                sid=submission_obj.sid,
                state='submitted'
            ))
            self.datastore.submission.save(sub.sid, sub)

            self.log.debug("Submission complete. Dispatching: %s", sub.sid)
            self.dispatcher.dispatch_submission(sub, completed_queue=completed_queue)

            return sub
        finally:
            # Just in case this method fails clean up local files
            if cleanup:
                for path in local_files:
                    if path and os.path.exists(path):
                        # noinspection PyBroadException
                        try:
                            os.unlink(path)
                        except Exception:
                            self.log.error("Couldn't delete dangling file %s", path)