lib/cuckoo/core/database.py
# Copyright (C) 2010-2012 Cuckoo Sandbox Developers.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.
import os
import sys
import MySQLdb
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, relationship
from sqlalchemy.sql.expression import desc
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy import Column, Integer, Text, DateTime, ForeignKey
from lib.cuckoo.common.constants import CUCKOO_ROOT
from lib.cuckoo.common.exceptions import CuckooDatabaseError, CuckooOperationalError
from lib.cuckoo.common.abstracts import Dictionary
from lib.cuckoo.common.utils import create_folder
# SQL Alchemy init time
engine = create_engine("mysql://avtest:avtest@10.0.20.1:3306/avtest")
Session = sessionmaker(bind=engine)
s = Session()
Base = declarative_base()
# Table Exe abstraction
class Exe(Base):
__tablename__ = "exe"
id = Column(Integer, primary_key=True)
file_path = Column(Text)
md5 = Column(Text)
def __init__(self, file_path, md5):
self.file_path = file_path
self.md5 = md5
# Table Analysis abstraction
class Analysis(Base):
__tablename__ = "analysis"
id = Column(Integer, primary_key=True)
desc = Column(Text)
exe_id = Column('exe_id', Integer, ForeignKey('exe.id'))
exe = relationship("Exe")
created_on = Column(DateTime)
completed_on = Column(DateTime)
lock = Column(Integer)
status = Column(Integer)
def __init__(self, desc, exe_id):
self.desc = desc
self.exe_id = exe_id
# Table Tasks abstraction
class Task(Base):
__tablename__ = "tasks"
id = Column(Integer, primary_key=True)
a_id = Column("anal_id", Integer, ForeignKey('analysis.id'))
md5 = Column(Text)
file_path = Column(Text)
timeout = Column(Integer)
priority = Column(Integer)
custom = Column(Text)
machine = Column(Text)
package = Column(Text)
options = Column(Text)
platform = Column(Text)
lock = Column(Integer)
status = Column(Integer)
detected = Column(Integer)
def __init__(self,
file_path,
a_id,
md5,
timeout,
package,
options,
priority,
custom,
machine,
platform):
self.a_id = a_id
self.md5 = md5
self.file_path = file_path
self.timeout = timeout
self.priority = priority
self.custom = custom
self.machine = machine
self.package = package
self.options = options
self.platform = platform
self.lock = 0
self.status = 0
class Database:
"""Analysis queue database."""
def generate(self):
"""Create database.
@return: operation status.
NOTE: you need to create database and user before generate db tables
"""
try:
Base.metadata.create_all(engine)
except SQLAlchemyError as e:
raise CuckooDatabaseError("Unable to create database: %s" % e)
return True
'''
def generate(self):
conn = MySQLdb.connect("10.0.20.1", "avtest", "avtest", "avtest")
cursor = conn.cursor()
try:
"""We need 3 tables: Analysis, tasks and another one
needed for executables
"""
cursor.execute("CREATE TABLE analysis (" \
" `id` INTEGER NOT NULL AUTO_INCREMENT PRIMARY KEY," \
" `desc` TEXT NULL," \
" `exe_id` INTEGER NOT NULL," \
" `created_on` TIMESTAMP DEFAULT CURRENT_TIMESTAMP, " \
" `completed_on` TIMESTAMP NULL, " \
" `lock` INTEGER DEFAULT 0, " \
# Status possible values:
# 0 = not completed
# 1 = error occurred
# 2 = completed successfully.
" `status` INTEGER DEFAULT 0" \
");")
cursor.execute("CREATE TABLE exe (" \
" `id` INTEGER NOT NULL AUTO_INCREMENT PRIMARY KEY," \
" `file_path` TEXT NOT NULL, " \
" `md5` TEXT NULL " \
");")
cursor.execute("CREATE TABLE tasks ( " \
" `id` INTEGER NOT NULL AUTO_INCREMENT PRIMARY KEY, " \
" `anal_id` INTEGER NOT NULL, " \
" `md5` TEXT NULL, " \
" `file_path` TEXT NOT NULL, " \
" `timeout` INTEGER DEFAULT NULL, " \
" `priority` INTEGER DEFAULT 0, " \
" `custom` TEXT NULL, " \
" `machine` TEXT NULL, " \
" `package` TEXT NULL, " \
" `options` TEXT NULL, " \
" `platform` TEXT NULL, " \
" `added_on` TIMESTAMP DEFAULT CURRENT_TIMESTAMP, " \
" `completed_on` TIMESTAMP NULL, " \
" `lock` INTEGER DEFAULT 0, " \
# Status possible values:
# 0 = not completed
# 1 = error occurred
# 2 = completed successfully.
" `status` INTEGER DEFAULT 0, " \
# Detected possible values:
# 0 = not completed
# 1 = not detected (success)
# 2 = detected (fail!)
" `detected` INTEGER DEFAULT 0 " \
");")
except MySQLdb.Error as e:
raise CuckooDatabaseError("Unable to create database: %s" % e)
return True
'''
def add(self,
file_path,
a_id,
md5="",
timeout=0,
package="",
options="",
priority=1,
custom="",
machine="",
platform=""):
"""Add a task to database.
@param file_path: sample path.
@param anal_id: analysis id referenced to table.
@param md5: sample MD5.
@param timeout: selected timeout.
@param options: analysis options.
@param priority: analysis priority.
@param custom: custom options.
@param machine: selected machine.
@param platform: platform
@return: cursor or None.
"""
if not file_path or not os.path.exists(file_path):
return None
if not timeout:
timeout = 0
try:
task = Task(file_path,
a_id,
md5,
timeout,
package,
options,
priority,
custom,
machine,
platform)
s.add(task)
s.commit()
return task.id
except MySQLdb.Error as e:
raise CuckooDatabaseError("Unable to add task: %s" % e)
def add_analysis(self, desc, exe_id):
""" Add an analysis on database
@param desc: description
@param exe_id: id of executable to test
@return: cursor or None
"""
a = Analysis(desc, exe_id)
try:
s.add(a)
s.commit()
return a.id
except SQLAlchemyError as e:
raise CuckooDatabaseError("Unable to create analysis: %s" % e)
def add_exe(self, file_path, md5):
""" Add an exe to db
@param file_path: path to file
@return: cursor or None
"""
if not file_path or not os.path.exists(file_path):
return None
row = s.query(Exe).filter_by(md5=md5).first()
if row is not None:
return row.id
try:
exe = Exe(file_path, md5)
s.add(exe)
s.commit()
return exe.id
except SQLAlchemyError as e:
raise CuckooDatabaseError("Unable to add executable, reason: %s" % e)
def fetch(self):
try:
s = Session()
task = s.query(Task).filter_by(lock=0, status=0).order_by(desc(Task.priority)).first()
#print("fetching task %s" % task)
s.close()
return task
except SQLAlchemyError as e:
raise CuckooDatabaseError("Unable to fetch, reason: %s" % e)
def lock(self, task_id):
"""Lock a task.
@param task_id: task id.
@return: operation status.
"""
try:
task = s.query(Task).filter_by(id=task_id).first()
except SQLAlchemyError as e:
raise CuckooDatabaseError("Unable to find lock, reason: %s" % e)
if task:
try:
task.lock = 1
s.commit()
except SQLAlchemyError as e:
raise CuckooDatabaseError("Unable to update lock, reason: %s" % e)
else:
return False
return True
def unlock(self, task_id):
"""Unlock a task.
@param task_id: task id.
@return: operation status.
"""
try:
task = s.query(Task).filter_by(id=task_id).first()
except SQLAlchemyError as e:
raise CuckooDatabaseError("Unable to find lock, reason: %s" % e)
if row:
try:
task.lock = 0
s.commit()
except SQLAlchemyError as e:
raise CuckooDatabaseError("Unable to unlock, reason: %s" % e)
else:
return False
return True
def complete(self, task_id, success=True):
"""Mark a task as completed.
@param task_id: task id.
@param success: completed with status.
@return: operation status.
"""
try:
task = s.query(Task).filter_by(id=task_id).first()
except SQLAlchemyError as e:
raise CuckooDatabaseError("Unable to find locked task, reason: %s" % e)
if task:
if success:
task.status = 2
else:
task.status = 1
try:
s.commit()
except SQLAlchemyError as e:
raise CuckooDatabaseError("Unable to complete, reason: %s" % e)
else:
return False
return True
def get_all_analysis(self):
""" Get all analysis
@return ALL analysis
"""
try:
analysis = s.query(Analysis).order_by(desc(Analysis.status),desc(Analysis.created_on)).all()
return analysis
except SQLAlchemyError as e:
raise CuckooDatabaseError("Unable to get all analysis, reason:" % e)
def get_analysis(self, a_id):
""" Retrive all tasks for given analysis id
@param: a_id: analysis id
@returns: all tasks for that analysis
"""
try:
s = Session()
tasks = s.query(Task).filter_by(a_id=a_id).all()
return tasks
except SQLAlchemyError as e:
raise CuckooDatabaseError("Unable to get all tasks for analysis, reason: %s" % e)
def get_task(self, task_id):
try:
s = Session()
task = s.query(Analysis).filter_by(task_id=task_id).first()
return task
except SQLAlchemyError as e:
raise CuckooDatabaseError("Unable to get all tasks for analysis, reason: %s" % e)