Python Database.fetch方法代码示例

# 需要导入模块: from lib.cuckoo.core.database import Database [as 别名]
# 或者: from lib.cuckoo.core.database.Database import fetch [as 别名]

#.........这里部分代码省略.........

    def stop(self):
        """Stop scheduler."""
        self.running = False
        # Shutdown machine manager (used to kill machines that still alive).
        machinery.shutdown()

    def start(self):
        """Start scheduler."""
        self.initialize()

        log.info("Waiting for analysis tasks.")

        # Message queue with threads to transmit exceptions (used as IPC).
        errors = Queue.Queue()

        # Command-line overrides the configuration file.
        if self.maxcount is None:
            self.maxcount = self.cfg.cuckoo.max_analysis_count

        # This loop runs forever.
        while self.running:
            time.sleep(1)

            # Wait until the machine lock is not locked. This is only the case
            # when all machines are fully running, rather that about to start
            # or still busy starting. This way we won't have race conditions
            # with finding out there are no available machines in the analysis
            # manager or having two analyses pick the same machine.
            if not machine_lock.acquire(False):
                continue

            machine_lock.release()

            # If not enough free disk space is available, then we print an
            # error message and wait another round (this check is ignored
            # when the freespace configuration variable is set to zero).
            if self.cfg.cuckoo.freespace:
                # Resolve the full base path to the analysis folder, just in
                # case somebody decides to make a symbolic link out of it.
                dir_path = os.path.join(CUCKOO_ROOT, "storage", "analyses")

                # TODO: Windows support
                if hasattr(os, "statvfs"):
                    dir_stats = os.statvfs(dir_path)

                    # Calculate the free disk space in megabytes.
                    space_available = dir_stats.f_bavail * dir_stats.f_frsize
                    space_available /= 1024 * 1024

                    if space_available < self.cfg.cuckoo.freespace:
                        log.error("Not enough free disk space! (Only %d MB!)",
                                  space_available)
                        continue

            # Have we limited the number of concurrently executing machines?
            if self.cfg.cuckoo.max_machines_count > 0:
                # Are too many running?
                if len(machinery.running()) >= self.cfg.cuckoo.max_machines_count:
                    continue

            # If no machines are available, it's pointless to fetch for
            # pending tasks. Loop over.
            if not machinery.availables():
                continue

            # Exits if max_analysis_count is defined in the configuration
            # file and has been reached.
            if self.maxcount and self.total_analysis_count >= self.maxcount:
                if active_analysis_count <= 0:
                    log.debug("Reached max analysis count, exiting.")
                    self.stop()
                continue

            # Fetch a pending analysis task.
            # TODO This fixes only submissions by --machine, need to add
            # other attributes (tags etc).
            # TODO We should probably move the entire "acquire machine" logic
            # from the Analysis Manager to the Scheduler and then pass the
            # selected machine onto the Analysis Manager instance.
            for machine in self.db.get_available_machines():

                task = self.db.fetch(machine=machine.name)
                if task:
                    log.debug("Processing task #%s", task.id)
                    self.total_analysis_count += 1

                    # Initialize and start the analysis manager.
                    analysis = AnalysisManager(task, errors)
                    analysis.daemon = True
                    analysis.start()
                    break

            # Deal with errors.
            try:
                raise errors.get(block=False)
            except Queue.Empty:
                pass

        log.debug("End of analyses.")

# 需要导入模块: from lib.cuckoo.core.database import Database [as 别名]
# 或者: from lib.cuckoo.core.database.Database import fetch [as 别名]

#.........这里部分代码省略.........
        machinery = plugin()

        # Find its configuration file.
        conf = os.path.join(CUCKOO_ROOT, "conf", "%s.conf" % machinery_name)

        if not os.path.exists(conf):
            raise CuckooCriticalError("The configuration file for machine "
                                      "manager \"{0}\" does not exist at path:"
                                      " {1}".format(machinery_name, conf))

        # Provide a dictionary with the configuration options to the
        # machine manager instance.
        machinery.set_options(Config(conf))
        # Initialize the machine manager.
        try:
            machinery.initialize(machinery_name)
        except CuckooMachineError as e:
            raise CuckooCriticalError("Error initializing machines: %s" % e)

        # At this point all the available machines should have been identified
        # and added to the list. If none were found, Cuckoo needs to abort the
        # execution.
        if len(machinery.machines()) == 0:
            raise CuckooCriticalError("No machines available")
        else:
            log.info("Loaded %s machine/s", len(machinery.machines()))

    def stop(self):
        """Stop scheduler."""
        self.running = False
        # Shutdown machine manager (used to kill machines that still alive).
        machinery.shutdown()

    def start(self):
        """Start scheduler."""
        global total_analysis_count
        self.initialize()

        log.info("Waiting for analysis tasks...")

        # Message queue with threads to transmit exceptions (used as IPC).
        errors = Queue.Queue()

        maxcount = self.cfg.cuckoo.max_analysis_count

        # This loop runs forever.
        while self.running:
            time.sleep(1)

            # If not enough free diskspace is available, then we print an
            # error message and wait another round (this check is ignored
            # when freespace is set to zero).
            if self.cfg.cuckoo.freespace:
                # Resolve the full base path to the analysis folder, just in
                # case somebody decides to make a symlink out of it.
                dir_path = os.path.join(CUCKOO_ROOT, "storage", "analyses")

                # TODO: Windows support
                if hasattr(os, "statvfs"):
                    dir_stats = os.statvfs(dir_path)

                    # Free diskspace in megabytes.
                    space_available = dir_stats.f_bavail * dir_stats.f_frsize
                    space_available /= 1024 * 1024

                    if space_available < self.cfg.cuckoo.freespace:
                        log.error("Not enough free diskspace! (Only %d MB!)",
                                  space_available)
                        continue

            # If no machines are available, it's pointless to fetch for
            # pending tasks. Loop over.
            if machinery.availables() == 0:
                continue

            # Exits if max_analysis_count is defined in config file and
            # is reached.
            if maxcount and total_analysis_count >= maxcount:
                if active_analysis_count <= 0:
                    self.stop()
            else:
                # Fetch a pending analysis task.
                task = self.db.fetch()

                if task:
                    log.debug("Processing task #%s", task.id)
                    total_analysis_count += 1

                    # Initialize the analysis manager.
                    analysis = AnalysisManager(task, errors)
                    # Start.
                    analysis.start()

            # Deal with errors.
            try:
                error = errors.get(block=False)
            except Queue.Empty:
                pass
            else:
                raise error

# 需要导入模块: from lib.cuckoo.core.database import Database [as 别名]
# 或者: from lib.cuckoo.core.database.Database import fetch [as 别名]
class TestDatabase:
    def setUp(self):
        self.tmp = os.path.join(tempfile.mkdtemp(), "dbtestcuckoo")
        self.d = Database(db_file=self.tmp)

    def test_db_path_default(self):
        """@note: Regression unit test."""
        d = Database()
        assert_equals(d.db_file, os.path.join(CUCKOO_ROOT, "db", "cuckoo.db"))
        assert os.path.exists(self.d.db_file)

    def test_db_path_custom(self):
        """@note: Regression unit test."""
        tmp = tempfile.mkstemp()[1]
        d = Database(db_file=tmp)
        assert_equals(d.db_file, tmp)
        assert os.path.exists(self.d.db_file)
        os.remove(tmp)

    def test_generate(self):
        conn = sqlite3.connect(self.tmp)
        cursor = conn.cursor()
        cursor.execute("SELECT count(name) FROM sqlite_master WHERE name='tasks';")
        assert_equals(1, cursor.fetchone()[0])

    def test_add(self):
        tmp = tempfile.mkstemp()[1]
        assert_equals(1, self.d.add(file_path=tmp))
        conn = sqlite3.connect(self.tmp)
        cursor = conn.cursor()
        cursor.execute("SELECT count(*) FROM tasks;")
        assert_equals(1, cursor.fetchone()[0])
        os.remove(tmp)

    def test_add_file_not_found(self):
        assert_equals(None, self.d.add(file_path="foo"))

    def test_fetch(self):
        tmp = tempfile.mkstemp()[1]
        assert_equals(1, self.d.add(file_path=tmp))
        assert_equals(tmp, self.d.fetch()["file_path"])
        conn = sqlite3.connect(self.tmp)
        cursor = conn.cursor()
        cursor.execute("SELECT count(*) FROM tasks;")
        assert_equals(1, cursor.fetchone()[0])
        os.remove(tmp)

    def test_lock(self):
        tmp = tempfile.mkstemp()[1]
        assert_equals(1, self.d.add(file_path=tmp))
        assert self.d.lock(1)
        conn = sqlite3.connect(self.tmp)
        cursor = conn.cursor()
        cursor.execute("SELECT count(*) FROM tasks WHERE lock=0;")
        assert_equals(0, cursor.fetchone()[0])
        os.remove(tmp)

    def test_unlock(self):
        tmp = tempfile.mkstemp()[1]
        assert_equals(1, self.d.add(file_path=tmp))
        assert self.d.lock(1)
        conn = sqlite3.connect(self.tmp)
        cursor = conn.cursor()
        cursor.execute("SELECT count(*) FROM tasks WHERE lock=0;")
        assert_equals(0, cursor.fetchone()[0])
        assert self.d.unlock(1)
        cursor.execute("SELECT count(*) FROM tasks WHERE lock=0;")
        assert_equals(1, cursor.fetchone()[0])
        os.remove(tmp)

    def test_complete_success(self):
        tmp = tempfile.mkstemp()[1]
        assert_equals(1, self.d.add(file_path=tmp))
        assert self.d.complete(1, True)
        conn = sqlite3.connect(self.tmp)
        cursor = conn.cursor()
        cursor.execute("SELECT count(*) FROM tasks WHERE status=2;")
        assert_equals(1, cursor.fetchone()[0])
        os.remove(tmp)

    def test_complete_fail(self):
        tmp = tempfile.mkstemp()[1]
        assert_equals(1, self.d.add(file_path=tmp))
        assert self.d.complete(1, False)
        conn = sqlite3.connect(self.tmp)
        cursor = conn.cursor()
        cursor.execute("SELECT count(*) FROM tasks WHERE status=1;")
        assert_equals(1, cursor.fetchone()[0])
        os.remove(tmp)

    def tearDown(self):
        os.remove(self.tmp)

# 需要导入模块: from lib.cuckoo.core.database import Database [as 别名]
# 或者: from lib.cuckoo.core.database.Database import fetch [as 别名]
class Scheduler:
    """Task scheduler."""

    def __init__(self):
        self.running = True
        self.cfg = Config()
        self.db = Database()

    def initialize(self):
        """Initialize machine manager."""
        global mmanager

        log.info("Using \"%s\" machine manager" % self.cfg.cuckoo.machine_manager)
        name = "modules.machinemanagers.%s" % self.cfg.cuckoo.machine_manager

        try:
            __import__(name, globals(), locals(), ["dummy"], -1)
        except ImportError as e:
            raise CuckooMachineError("Unable to import machine manager plugin: %s" % e)

        MachineManager()
        module = MachineManager.__subclasses__()[0]
        mmanager = module()
        mmanager_conf = os.path.join(CUCKOO_ROOT, "conf", "%s.conf" % self.cfg.cuckoo.machine_manager)

        if not os.path.exists(mmanager_conf):
            raise CuckooMachineError("The configuration file for machine manager \"%s\" does not exist at path: %s"
                                     % (self.cfg.cuckoo.machine_manager, mmanager_conf))

        mmanager.set_options(Config(mmanager_conf))
        mmanager.initialize(self.cfg.cuckoo.machine_manager)

        if len(mmanager.machines) == 0:
            raise CuckooMachineError("No machines available")
        else:
            log.info("Loaded %s machine/s" % len(mmanager.machines))

    def stop(self):
        """Stop scheduler."""
        self.running = False

        # Shutdown vm alive.
        # TODO: in future this code may be moved.
        if len(mmanager.running()) > 0:
            log.info("Shutting down guests")
            for machine in mmanager.running():
                try:
                    mmanager.stop(machine.label)
                except CuckooMachineError as e:
                    log.error("Unable to shutdown machine %s, please check manually. Error: %s" % (machine.label, e.message))

    def start(self):
        """Start scheduler."""
        self.initialize()

        log.info("Waiting for analysis tasks...")

        while self.running:
            time.sleep(1)

            if mmanager.availables() == 0:
                log.debug("No machines available, try again")
                continue

            task = self.db.fetch()

            if not task:
                log.debug("No pending tasks, try again")
                continue

            analysis = AnalysisManager(task)
            analysis.daemon = True
            analysis.start()

# 需要导入模块: from lib.cuckoo.core.database import Database [as 别名]
# 或者: from lib.cuckoo.core.database.Database import fetch [as 别名]

#.........这里部分代码省略.........
        if len(machinery.machines()) > 1 and self.db.engine.name == "sqlite":
            log.warning("The SQLite database is not compatible with "
                        "multi-threaded use-cases such as running multiple "
                        "virtual machine in parallel. Please upgrade to "
                        "PostgreSQL or MySQL when running multiple VMs.")

        if len(machinery.machines()) > 3 and self.cfg.cuckoo.process_results:
            log.warning("When running many virtual machines it is recommended "
                        "to process the results in a separate process.py to "
                        "increase throughput and stability. Please read the "
                        "documentation about the `Processing Utility`.")

    def stop(self):
        """Stop scheduler."""
        self.running = False
        # Shutdown machine manager (used to kill machines that still alive).
        machinery.shutdown()

    def start(self):
        """Start scheduler."""
        self.initialize()

        log.info("Waiting for analysis tasks.")

        # Message queue with threads to transmit exceptions (used as IPC).
        errors = Queue.Queue()

        # Command-line overrides the configuration file.
        if self.maxcount is None:
            self.maxcount = self.cfg.cuckoo.max_analysis_count

        # This loop runs forever.
        while self.running:
            time.sleep(1)

            # Wait until the machine lock is not locked, that is, all machines
            # have been assigned a machine and are running.
            if not machine_lock.acquire(False):
                continue

            machine_lock.release()

            # If not enough free disk space is available, then we print an
            # error message and wait another round (this check is ignored
            # when the freespace configuration variable is set to zero).
            if self.cfg.cuckoo.freespace:
                # Resolve the full base path to the analysis folder, just in
                # case somebody decides to make a symbolic link out of it.
                dir_path = os.path.join(CUCKOO_ROOT, "storage", "analyses")

                # TODO: Windows support
                if hasattr(os, "statvfs"):
                    dir_stats = os.statvfs(dir_path)

                    # Calculate the free disk space in megabytes.
                    space_available = dir_stats.f_bavail * dir_stats.f_frsize
                    space_available /= 1024 * 1024

                    if space_available < self.cfg.cuckoo.freespace:
                        log.error("Not enough free disk space! (Only %d MB!)",
                                  space_available)
                        continue

            # Have we limited the number of concurrently executing machines?
            if self.cfg.cuckoo.max_machines_count > 0:
                # Are too many running?
                if len(machinery.running()) >= self.cfg.cuckoo.max_machines_count:
                    continue

            # Exits if max_analysis_count is defined in the configuration
            # file and has been reached.
            if self.maxcount and self.total_analysis_count >= self.maxcount:
                if active_analysis_count <= 0:
                    self.stop()
            else:
                task = None
                if machinery.availables():
                    # Machines available, get the next task
                    # lock the machine, and run the analysis.
                    task = self.db.fetch(status=TASK_PENDING)

                if not task:
                    # No pending task available, simply run scheduled task
                    task = self.db.fetch(status=TASK_SCHEDULED)

                if task:
                    log.debug("Processing task #%s", task.id)

                    self.total_analysis_count += 1

                    # Initialize and start the analysis manager.
                    analysis = AnalysisManager(task, errors)
                    analysis.start()
                    task = None

            # Deal with errors.
            try:
                raise errors.get(block=False)
            except Queue.Empty:
                pass

# 需要导入模块: from lib.cuckoo.core.database import Database [as 别名]
# 或者: from lib.cuckoo.core.database.Database import fetch [as 别名]
class Scheduler:
    """Task scheduler."""

    def __init__(self):
        self.running = True
        self.cfg = Config()
        self.db = Database()

    def initialize(self):
        """Initialize machine manager."""
        global mmanager

        log.info('Using "%s" machine manager' % self.cfg.cuckoo.machine_manager)
        name = "modules.machinemanagers.%s" % self.cfg.cuckoo.machine_manager

        try:
            __import__(name, globals(), locals(), ["dummy"], -1)
        except ImportError as e:
            raise CuckooMachineError("Unable to import machine manager plugin: %s" % e)

        MachineManager()
        module = MachineManager.__subclasses__()[0]
        mmanager = module()
        mmanager_conf = os.path.join(CUCKOO_ROOT, "conf", "%s.conf" % self.cfg.cuckoo.machine_manager)

        if not os.path.exists(mmanager_conf):
            raise CuckooMachineError(
                'The configuration file for machine manager "%s" does not exist at path: %s'
                % (self.cfg.cuckoo.machine_manager, mmanager_conf)
            )

        mmanager.set_options(Config(mmanager_conf))
        mmanager.initialize(self.cfg.cuckoo.machine_manager)

        if len(mmanager.machines) == 0:
            raise CuckooMachineError("No machines available")
        else:
            log.info("Loaded %s machine/s" % len(mmanager.machines))

    def stop(self):
        """Stop scheduler."""
        self.running = False

        # Shutdown machine manager (used to kill vm still alive).
        mmanager.shutdown()

    def start(self):
        """Start scheduler."""
        self.initialize()

        log.info("Waiting for analysis tasks...")

        while self.running:
            time.sleep(1)

            if mmanager.availables() == 0:
                # log.debug("No machines available, try again")
                continue
            db_lock.acquire()
            task = self.db.fetch()

            if task:
                log.debug("Locking db task #%s" % task.id)
                self.db.lock(task.id)
                db_lock.release()
                # Go with analysis.
                analysis = AnalysisManager(task)
                analysis.daemon = True
                analysis.start()
            else:
                # log.debug("No pending tasks, try again")
                db_lock.release()