Package lib :: Package cuckoo :: Package core :: Module database
[hide private]
[frames] | no frames]

Source Code for Module lib.cuckoo.core.database

   1  # Copyright (C) 2010-2014 Cuckoo Sandbox Developers. 
   2  # This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org 
   3  # See the file 'docs/LICENSE' for copying permission. 
   4   
   5  import os 
   6  import json 
   7  import logging 
   8  from datetime import datetime 
   9   
  10  from lib.cuckoo.common.config import Config 
  11  from lib.cuckoo.common.constants import CUCKOO_ROOT 
  12  from lib.cuckoo.common.exceptions import CuckooDatabaseError 
  13  from lib.cuckoo.common.exceptions import CuckooOperationalError 
  14  from lib.cuckoo.common.exceptions import CuckooDependencyError 
  15  from lib.cuckoo.common.objects import File, URL 
  16  from lib.cuckoo.common.utils import create_folder, Singleton 
  17   
  18  try: 
  19      from sqlalchemy import create_engine, Column 
  20      from sqlalchemy import Integer, String, Boolean, DateTime, Enum 
  21      from sqlalchemy import ForeignKey, Text, Index, Table 
  22      from sqlalchemy.ext.declarative import declarative_base 
  23      from sqlalchemy.exc import SQLAlchemyError, IntegrityError 
  24      from sqlalchemy.orm import sessionmaker, relationship, joinedload, backref 
  25      from sqlalchemy.pool import NullPool 
  26      Base = declarative_base() 
  27  except ImportError: 
  28      raise CuckooDependencyError("Unable to import sqlalchemy " 
  29                                  "(install with `pip install sqlalchemy`)") 
  30   
  31  log = logging.getLogger(__name__) 
  32   
  33  TASK_PENDING = "pending" 
  34  TASK_RUNNING = "running" 
  35  TASK_COMPLETED = "completed" 
  36  TASK_RECOVERED = "recovered" 
  37  TASK_REPORTED = "reported" 
  38  TASK_FAILED_ANALYSIS = "failed_analysis" 
  39  TASK_FAILED_PROCESSING = "failed_processing" 
  40   
  41  # Secondary table used in association Machine - Tag. 
  42  machines_tags = Table("machines_tags", Base.metadata, 
  43      Column("machine_id", Integer, ForeignKey("machines.id")), 
  44      Column("tag_id", Integer, ForeignKey("tags.id")) 
  45  ) 
  46   
  47  # Secondary table used in association Task - Tag. 
  48  tasks_tags = Table("tasks_tags", Base.metadata, 
  49      Column("task_id", Integer, ForeignKey("tasks.id")), 
  50      Column("tag_id", Integer, ForeignKey("tags.id")) 
  51  ) 
  52   
53 -class Machine(Base):
54 """Configured virtual machines to be used as guests.""" 55 __tablename__ = "machines" 56 57 id = Column(Integer(), primary_key=True) 58 name = Column(String(255), nullable=False) 59 label = Column(String(255), nullable=False) 60 ip = Column(String(255), nullable=False) 61 platform = Column(String(255), nullable=False) 62 tags = relationship("Tag", secondary=machines_tags, cascade="all, delete", 63 single_parent=True, backref=backref("machine", cascade="all")) 64 interface = Column(String(255), nullable=True) 65 snapshot = Column(String(255), nullable=True) 66 locked = Column(Boolean(), nullable=False, default=False) 67 locked_changed_on = Column(DateTime(timezone=False), nullable=True) 68 status = Column(String(255), nullable=True) 69 status_changed_on = Column(DateTime(timezone=False), nullable=True) 70 resultserver_ip = Column(String(255), nullable=False) 71 resultserver_port = Column(String(255), nullable=False) 72
73 - def __repr__(self):
74 return "<Machine('{0}','{1}')>".format(self.id, self.name)
75
76 - def to_dict(self):
77 """Converts object to dict. 78 @return: dict 79 """ 80 d = {} 81 for column in self.__table__.columns: 82 value = getattr(self, column.name) 83 if isinstance(value, datetime): 84 d[column.name] = value.strftime("%Y-%m-%d %H:%M:%S") 85 else: 86 d[column.name] = value 87 88 # Tags are a relation so no column to iterate. 89 d["tags"] = [tag.name for tag in self.tags] 90 91 return d
92
93 - def to_json(self):
94 """Converts object to JSON. 95 @return: JSON data 96 """ 97 return json.dumps(self.to_dict())
98
99 - def __init__(self, 100 name, 101 label, 102 ip, 103 platform, 104 interface, 105 snapshot, 106 resultserver_ip, 107 resultserver_port):
108 self.name = name 109 self.label = label 110 self.ip = ip 111 self.platform = platform 112 self.interface = interface 113 self.snapshot = snapshot 114 self.resultserver_ip = resultserver_ip 115 self.resultserver_port = resultserver_port
116
117 -class Tag(Base):
118 """Tag describing anything you want.""" 119 __tablename__ = "tags" 120 121 id = Column(Integer(), primary_key=True) 122 name = Column(String(255), nullable=False, unique=True) 123
124 - def __repr__(self):
125 return "<Tag('{0}','{1}')>".format(self.id, self.name)
126
127 - def __init__(self, 128 name):
129 self.name = name
130
131 -class Guest(Base):
132 """Tracks guest run.""" 133 __tablename__ = "guests" 134 135 id = Column(Integer(), primary_key=True) 136 name = Column(String(255), nullable=False) 137 label = Column(String(255), nullable=False) 138 manager = Column(String(255), nullable=False) 139 started_on = Column(DateTime(timezone=False), 140 default=datetime.now, 141 nullable=False) 142 shutdown_on = Column(DateTime(timezone=False), nullable=True) 143 task_id = Column(Integer, 144 ForeignKey("tasks.id"), 145 nullable=False, 146 unique=True) 147
148 - def __repr__(self):
149 return "<Guest('{0}','{1}')>".format(self.id, self.name)
150
151 - def to_dict(self):
152 """Converts object to dict. 153 @return: dict 154 """ 155 d = {} 156 for column in self.__table__.columns: 157 value = getattr(self, column.name) 158 if isinstance(value, datetime): 159 d[column.name] = value.strftime("%Y-%m-%d %H:%M:%S") 160 else: 161 d[column.name] = value 162 return d
163
164 - def to_json(self):
165 """Converts object to JSON. 166 @return: JSON data 167 """ 168 return json.dumps(self.to_dict())
169
170 - def __init__(self, name, label, manager):
171 self.name = name 172 self.label = label 173 self.manager = manager
174
175 -class Sample(Base):
176 """Submitted files details.""" 177 __tablename__ = "samples" 178 179 id = Column(Integer(), primary_key=True) 180 file_size = Column(Integer(), nullable=False) 181 file_type = Column(String(255), nullable=False) 182 md5 = Column(String(32), nullable=False) 183 crc32 = Column(String(8), nullable=False) 184 sha1 = Column(String(40), nullable=False) 185 sha256 = Column(String(64), nullable=False) 186 sha512 = Column(String(128), nullable=False) 187 ssdeep = Column(String(255), nullable=True) 188 __table_args__ = (Index("hash_index", 189 "md5", 190 "crc32", 191 "sha1", 192 "sha256", 193 "sha512", 194 unique=True), ) 195
196 - def __repr__(self):
197 return "<Sample('{0}','{1}')>".format(self.id, self.sha256)
198
199 - def to_dict(self):
200 """Converts object to dict. 201 @return: dict 202 """ 203 d = {} 204 for column in self.__table__.columns: 205 value = getattr(self, column.name) 206 d[column.name] = value 207 return d
208
209 - def to_json(self):
210 """Converts object to JSON. 211 @return: JSON data 212 """ 213 return json.dumps(self.to_dict())
214
215 - def __init__(self, 216 md5, 217 crc32, 218 sha1, 219 sha256, 220 sha512, 221 file_size, 222 file_type=None, 223 ssdeep=None):
224 self.md5 = md5 225 self.sha1 = sha1 226 self.crc32 = crc32 227 self.sha256 = sha256 228 self.sha512 = sha512 229 self.file_size = file_size 230 if file_type: 231 self.file_type = file_type 232 if ssdeep: 233 self.ssdeep = ssdeep
234
235 -class Error(Base):
236 """Analysis errors.""" 237 __tablename__ = "errors" 238 239 id = Column(Integer(), primary_key=True) 240 message = Column(String(255), nullable=False) 241 task_id = Column(Integer, ForeignKey("tasks.id"), nullable=False) 242
243 - def to_dict(self):
244 """Converts object to dict. 245 @return: dict 246 """ 247 d = {} 248 for column in self.__table__.columns: 249 value = getattr(self, column.name) 250 d[column.name] = value 251 return d
252
253 - def to_json(self):
254 """Converts object to JSON. 255 @return: JSON data 256 """ 257 return json.dumps(self.to_dict())
258
259 - def __init__(self, message, task_id):
260 self.message = message 261 self.task_id = task_id
262
263 - def __repr__(self):
264 return "<Error('{0}','{1}','{2}')>".format(self.id, self.message, self.task_id)
265
266 -class Task(Base):
267 """Analysis task queue.""" 268 __tablename__ = "tasks" 269 270 id = Column(Integer(), primary_key=True) 271 target = Column(Text(), nullable=False) 272 category = Column(String(255), nullable=False) 273 timeout = Column(Integer(), server_default="0", nullable=False) 274 priority = Column(Integer(), server_default="1", nullable=False) 275 custom = Column(String(255), nullable=True) 276 machine = Column(String(255), nullable=True) 277 package = Column(String(255), nullable=True) 278 tags = relationship("Tag", secondary=tasks_tags, cascade="all, delete", 279 single_parent=True, backref=backref("task", cascade="all"), 280 lazy="subquery") 281 options = Column(String(255), nullable=True) 282 platform = Column(String(255), nullable=True) 283 memory = Column(Boolean, nullable=False, default=False) 284 enforce_timeout = Column(Boolean, nullable=False, default=False) 285 clock = Column(DateTime(timezone=False), 286 default=datetime.now, 287 nullable=False) 288 added_on = Column(DateTime(timezone=False), 289 default=datetime.now, 290 nullable=False) 291 started_on = Column(DateTime(timezone=False), nullable=True) 292 completed_on = Column(DateTime(timezone=False), nullable=True) 293 status = Column(Enum(TASK_PENDING, 294 TASK_RUNNING, 295 TASK_COMPLETED, 296 TASK_REPORTED, 297 TASK_RECOVERED, 298 name="status_type"), 299 server_default=TASK_PENDING, 300 nullable=False) 301 sample_id = Column(Integer, ForeignKey("samples.id"), nullable=True) 302 sample = relationship("Sample", backref="tasks") 303 guest = relationship("Guest", uselist=False, backref="tasks", cascade="save-update, delete") 304 errors = relationship("Error", backref="tasks", cascade="save-update, delete") 305
306 - def to_dict(self):
307 """Converts object to dict. 308 @return: dict 309 """ 310 d = {} 311 for column in self.__table__.columns: 312 value = getattr(self, column.name) 313 if isinstance(value, datetime): 314 d[column.name] = value.strftime("%Y-%m-%d %H:%M:%S") 315 else: 316 d[column.name] = value 317 318 # Tags are a relation so no column to iterate. 319 d["tags"] = [tag.name for tag in self.tags] 320 321 return d
322
323 - def to_json(self):
324 """Converts object to JSON. 325 @return: JSON data 326 """ 327 return json.dumps(self.to_dict())
328
329 - def __init__(self, target=None):
330 self.target = target
331
332 - def __repr__(self):
333 return "<Task('{0}','{1}')>".format(self.id, self.target)
334
335 -class Database(object):
336 """Analysis queue database. 337 338 This class handles the creation of the database user for internal queue 339 management. It also provides some functions for interacting with it. 340 """ 341 __metaclass__ = Singleton 342
343 - def __init__(self, dsn=None):
344 """@param dsn: database connection string.""" 345 cfg = Config() 346 347 if dsn: 348 self.engine = create_engine(dsn, poolclass=NullPool) 349 elif cfg.database.connection: 350 self.engine = create_engine(cfg.database.connection, poolclass=NullPool) 351 else: 352 db_file = os.path.join(CUCKOO_ROOT, "db", "cuckoo.db") 353 if not os.path.exists(db_file): 354 db_dir = os.path.dirname(db_file) 355 if not os.path.exists(db_dir): 356 try: 357 create_folder(folder=db_dir) 358 except CuckooOperationalError as e: 359 raise CuckooDatabaseError("Unable to create database directory: {0}".format(e)) 360 361 self.engine = create_engine("sqlite:///{0}".format(db_file), poolclass=NullPool) 362 363 # Disable SQL logging. Turn it on for debugging. 364 self.engine.echo = False 365 # Connection timeout. 366 if cfg.database.timeout: 367 self.engine.pool_timeout = cfg.database.timeout 368 else: 369 self.engine.pool_timeout = 60 370 # Create schema. 371 try: 372 Base.metadata.create_all(self.engine) 373 except SQLAlchemyError as e: 374 raise CuckooDatabaseError("Unable to create or connect to database: {0}".format(e)) 375 376 # Get db session. 377 self.Session = sessionmaker(bind=self.engine)
378
379 - def __del__(self):
380 """Disconnects pool.""" 381 self.engine.dispose()
382
383 - def _get_or_create(self, session, model, **kwargs):
384 """Get an ORM instance or create it if not exist. 385 @param session: SQLAlchemy session object 386 @param model: model to query 387 @return: row instance 388 """ 389 instance = session.query(model).filter_by(**kwargs).first() 390 if instance: 391 return instance 392 else: 393 instance = model(**kwargs) 394 return instance
395
396 - def clean_machines(self):
397 """Clean old stored machines and related tables.""" 398 # Secondary table. 399 # TODO: this is better done via cascade delete. 400 self.engine.execute(machines_tags.delete()) 401 402 session = self.Session() 403 try: 404 session.query(Machine).delete() 405 session.commit() 406 except SQLAlchemyError as e: 407 log.debug("Database error cleaning machines: {0}".format(e)) 408 session.rollback() 409 finally: 410 session.close()
411
412 - def add_machine(self, 413 name, 414 label, 415 ip, 416 platform, 417 tags, 418 interface, 419 snapshot, 420 resultserver_ip, 421 resultserver_port):
422 """Add a guest machine. 423 @param name: machine id 424 @param label: machine label 425 @param ip: machine IP address 426 @param platform: machine supported platform 427 @param interface: sniffing interface for this machine 428 @param snapshot: snapshot name to use instead of the current one, if configured 429 @param resultserver_ip: IP address of the Result Server 430 @param resultserver_port: port of the Result Server 431 """ 432 session = self.Session() 433 machine = Machine(name=name, 434 label=label, 435 ip=ip, 436 platform=platform, 437 interface=interface, 438 snapshot=snapshot, 439 resultserver_ip=resultserver_ip, 440 resultserver_port=resultserver_port) 441 # Deal with tags format (i.e. foo,bar,baz) 442 if tags: 443 for tag in tags.replace(" ", "").split(","): 444 machine.tags.append(self._get_or_create(session, Tag, name=tag)) 445 session.add(machine) 446 447 try: 448 session.commit() 449 except SQLAlchemyError as e: 450 log.debug("Database error adding machine: {0}".format(e)) 451 session.rollback() 452 finally: 453 session.close()
454
455 - def set_status(self, task_id, status):
456 """Set task status. 457 @param task_id: task identifier 458 @param status: status string 459 @return: operation status 460 """ 461 session = self.Session() 462 try: 463 row = session.query(Task).get(task_id) 464 row.status = status 465 466 if status == TASK_RUNNING: 467 row.started_on = datetime.now() 468 elif status == TASK_COMPLETED: 469 row.completed_on = datetime.now() 470 471 session.commit() 472 except SQLAlchemyError as e: 473 log.debug("Database error setting status: {0}".format(e)) 474 session.rollback() 475 finally: 476 session.close()
477
478 - def fetch(self, lock=True):
479 """Fetches a task waiting to be processed and locks it for running. 480 @return: None or task 481 """ 482 session = self.Session() 483 row = None 484 485 try: 486 row = session.query(Task).filter(Task.status == TASK_PENDING).order_by("priority desc, added_on").first() 487 488 if not row: 489 return None 490 491 if lock: 492 self.set_status(task_id=row.id, status=TASK_RUNNING) 493 session.refresh(row) 494 except SQLAlchemyError as e: 495 log.debug("Database error fetching task: {0}".format(e)) 496 session.rollback() 497 finally: 498 session.close() 499 500 return row
501
502 - def guest_start(self, task_id, name, label, manager):
503 """Logs guest start. 504 @param task_id: task identifier 505 @param name: vm name 506 @param label: vm label 507 @param manager: vm manager 508 @return: guest row id 509 """ 510 session = self.Session() 511 guest = Guest(name, label, manager) 512 try: 513 session.query(Task).get(task_id).guest = guest 514 session.commit() 515 session.refresh(guest) 516 except SQLAlchemyError as e: 517 log.debug("Database error logging guest start: {0}".format(e)) 518 session.rollback() 519 return None 520 finally: 521 session.close() 522 return guest.id
523
524 - def guest_remove(self, guest_id):
525 """Removes a guest start entry.""" 526 session = self.Session() 527 try: 528 guest = session.query(Guest).get(guest_id) 529 session.delete(guest) 530 session.commit() 531 except SQLAlchemyError as e: 532 log.debug("Database error logging guest remove: {0}".format(e)) 533 session.rollback() 534 return None 535 finally: 536 session.close()
537
538 - def guest_stop(self, guest_id):
539 """Logs guest stop. 540 @param guest_id: guest log entry id 541 """ 542 session = self.Session() 543 try: 544 session.query(Guest).get(guest_id).shutdown_on = datetime.now() 545 session.commit() 546 except SQLAlchemyError as e: 547 log.debug("Database error logging guest stop: {0}".format(e)) 548 session.rollback() 549 finally: 550 session.close()
551
552 - def list_machines(self, locked=False):
553 """Lists virtual machines. 554 @return: list of virtual machines 555 """ 556 session = self.Session() 557 try: 558 if locked: 559 machines = session.query(Machine).options(joinedload("tags")).filter(Machine.locked == True).all() 560 else: 561 machines = session.query(Machine).options(joinedload("tags")).all() 562 except SQLAlchemyError as e: 563 log.debug("Database error listing machines: {0}".format(e)) 564 return None 565 finally: 566 session.close() 567 return machines
568
569 - def lock_machine(self, name=None, platform=None, tags=None):
570 """Places a lock on a free virtual machine. 571 @param name: optional virtual machine name 572 @param platform: optional virtual machine platform 573 @param tags: optional tags required (list) 574 @return: locked machine 575 """ 576 session = self.Session() 577 578 # Preventive checks. 579 if name and platform: 580 # Wrong usage. 581 log.error("You can select machine only by name or by platform.") 582 return None 583 elif name and tags: 584 # Also wrong usage 585 log.error("You can select machine only by name or by tags.") 586 return None 587 588 try: 589 machines = session.query(Machine) 590 if name: 591 machines = machines.filter(Machine.name == name) 592 if platform: 593 machines = machines.filter(Machine.platform == platform) 594 if tags: 595 for tag in tags: 596 machines = machines.filter(Machine.tags.any(name=tag.name)) 597 598 # Check if there are any machines that satisfy the 599 # selection requirements. 600 if machines.count() == 0: 601 raise CuckooOperationalError("No machines match selection criteria") 602 603 # Get only free machines. 604 machines = machines.filter(Machine.locked == False) 605 # Get only one. 606 machine = machines.first() 607 except SQLAlchemyError as e: 608 log.debug("Database error locking machine: {0}".format(e)) 609 session.close() 610 return None 611 612 if machine: 613 machine.locked = True 614 machine.locked_changed_on = datetime.now() 615 try: 616 session.commit() 617 session.refresh(machine) 618 except SQLAlchemyError as e: 619 log.debug("Database error locking machine: {0}".format(e)) 620 session.rollback() 621 return None 622 finally: 623 session.close() 624 625 return machine
626
627 - def unlock_machine(self, label):
628 """Remove lock form a virtual machine. 629 @param label: virtual machine label 630 @return: unlocked machine 631 """ 632 session = self.Session() 633 try: 634 machine = session.query(Machine).filter(Machine.label == label).first() 635 except SQLAlchemyError as e: 636 log.debug("Database error unlocking machine: {0}".format(e)) 637 session.close() 638 return None 639 640 if machine: 641 machine.locked = False 642 machine.locked_changed_on = datetime.now() 643 try: 644 session.commit() 645 session.refresh(machine) 646 except SQLAlchemyError as e: 647 log.debug("Database error locking machine: {0}".format(e)) 648 session.rollback() 649 return None 650 finally: 651 session.close() 652 653 return machine
654
655 - def count_machines_available(self):
656 """How many virtual machines are ready for analysis. 657 @return: free virtual machines count 658 """ 659 session = self.Session() 660 try: 661 machines_count = session.query(Machine).filter(Machine.locked == False).count() 662 except SQLAlchemyError as e: 663 log.debug("Database error counting machines: {0}".format(e)) 664 return 0 665 finally: 666 session.close() 667 return machines_count
668
669 - def set_machine_status(self, label, status):
670 """Set status for a virtual machine. 671 @param label: virtual machine label 672 @param status: new virtual machine status 673 """ 674 session = self.Session() 675 try: 676 machine = session.query(Machine).filter(Machine.label == label).first() 677 except SQLAlchemyError as e: 678 log.debug("Database error setting machine status: {0}".format(e)) 679 session.close() 680 return 681 682 if machine: 683 machine.status = status 684 machine.status_changed_on = datetime.now() 685 try: 686 session.commit() 687 session.refresh(machine) 688 except SQLAlchemyError as e: 689 log.debug("Database error setting machine status: {0}".format(e)) 690 session.rollback() 691 finally: 692 session.close() 693 else: 694 session.close()
695
696 - def add_error(self, message, task_id):
697 """Add an error related to a task. 698 @param message: error message 699 @param task_id: ID of the related task 700 """ 701 session = self.Session() 702 error = Error(message=message, task_id=task_id) 703 session.add(error) 704 try: 705 session.commit() 706 except SQLAlchemyError as e: 707 log.debug("Database error adding error log: {0}".format(e)) 708 session.rollback() 709 finally: 710 session.close()
711 712 # The following functions are mostly used by external utils. 713
714 - def add(self, 715 obj, 716 timeout=0, 717 package="", 718 options="", 719 priority=1, 720 custom="", 721 machine="", 722 platform="", 723 tags=None, 724 memory=False, 725 enforce_timeout=False, 726 clock=None):
727 """Add a task to database. 728 @param obj: object to add (File or URL). 729 @param timeout: selected timeout. 730 @param options: analysis options. 731 @param priority: analysis priority. 732 @param custom: custom options. 733 @param machine: selected machine. 734 @param platform: platform. 735 @param tags: optional tags that must be set for machine selection 736 @param memory: toggle full memory dump. 737 @param enforce_timeout: toggle full timeout execution. 738 @param clock: virtual machine clock time 739 @return: cursor or None. 740 """ 741 session = self.Session() 742 743 # Convert empty strings and None values to a valid int 744 if not timeout: 745 timeout = 0 746 if not priority: 747 priority = 1 748 749 if isinstance(obj, File): 750 sample = Sample(md5=obj.get_md5(), 751 crc32=obj.get_crc32(), 752 sha1=obj.get_sha1(), 753 sha256=obj.get_sha256(), 754 sha512=obj.get_sha512(), 755 file_size=obj.get_size(), 756 file_type=obj.get_type(), 757 ssdeep=obj.get_ssdeep()) 758 session.add(sample) 759 760 try: 761 session.commit() 762 except IntegrityError: 763 session.rollback() 764 try: 765 sample = session.query(Sample).filter(Sample.md5 == obj.get_md5()).first() 766 except SQLAlchemyError: 767 session.close() 768 return None 769 except SQLAlchemyError as e: 770 log.debug("Database error adding task: {0}".format(e)) 771 session.close() 772 return None 773 774 task = Task(obj.file_path) 775 task.sample_id = sample.id 776 elif isinstance(obj, URL): 777 task = Task(obj.url) 778 779 task.category = obj.__class__.__name__.lower() 780 task.timeout = timeout 781 task.package = package 782 task.options = options 783 task.priority = priority 784 task.custom = custom 785 task.machine = machine 786 task.platform = platform 787 task.memory = memory 788 task.enforce_timeout = enforce_timeout 789 790 # Deal with tags format (i.e. foo,bar,baz) 791 if tags: 792 for tag in tags.replace(" ","").split(","): 793 task.tags.append(self._get_or_create(session, Tag, name=tag)) 794 795 if clock: 796 if isinstance(clock, str) or isinstance(clock, unicode): 797 try: 798 task.clock = datetime.strptime(clock, "%m-%d-%Y %H:%M:%S") 799 except ValueError: 800 log.warning("The date you specified has an invalid format, using current timestamp") 801 task.clock = datetime.now() 802 else: 803 task.clock = clock 804 805 session.add(task) 806 807 try: 808 session.commit() 809 task_id = task.id 810 except SQLAlchemyError as e: 811 log.debug("Database error adding task: {0}".format(e)) 812 session.rollback() 813 return None 814 finally: 815 session.close() 816 817 return task_id
818
819 - def add_path(self, 820 file_path, 821 timeout=0, 822 package="", 823 options="", 824 priority=1, 825 custom="", 826 machine="", 827 platform="", 828 tags=None, 829 memory=False, 830 enforce_timeout=False, 831 clock=None):
832 """Add a task to database from file path. 833 @param file_path: sample path. 834 @param timeout: selected timeout. 835 @param options: analysis options. 836 @param priority: analysis priority. 837 @param custom: custom options. 838 @param machine: selected machine. 839 @param platform: platform. 840 @param tags: Tags required in machine selection 841 @param memory: toggle full memory dump. 842 @param enforce_timeout: toggle full timeout execution. 843 @param clock: virtual machine clock time 844 @return: cursor or None. 845 """ 846 if not file_path or not os.path.exists(file_path): 847 return None 848 849 # Convert empty strings and None values to a valid int 850 if not timeout: 851 timeout = 0 852 if not priority: 853 priority = 1 854 855 return self.add(File(file_path), 856 timeout, 857 package, 858 options, 859 priority, 860 custom, 861 machine, 862 platform, 863 tags, 864 memory, 865 enforce_timeout, 866 clock)
867
868 - def add_url(self, 869 url, 870 timeout=0, 871 package="", 872 options="", 873 priority=1, 874 custom="", 875 machine="", 876 platform="", 877 tags=None, 878 memory=False, 879 enforce_timeout=False, 880 clock=None):
881 """Add a task to database from url. 882 @param url: url. 883 @param timeout: selected timeout. 884 @param options: analysis options. 885 @param priority: analysis priority. 886 @param custom: custom options. 887 @param machine: selected machine. 888 @param platform: platform. 889 @param tags: tags for machine selection 890 @param memory: toggle full memory dump. 891 @param enforce_timeout: toggle full timeout execution. 892 @param clock: virtual machine clock time 893 @return: cursor or None. 894 """ 895 896 # Convert empty strings and None values to a valid int 897 if not timeout: 898 timeout = 0 899 if not priority: 900 priority = 1 901 902 return self.add(URL(url), 903 timeout, 904 package, 905 options, 906 priority, 907 custom, 908 machine, 909 platform, 910 tags, 911 memory, 912 enforce_timeout, 913 clock)
914
915 - def reschedule(self, task_id):
916 """Reschedule a task. 917 @param task_id: ID of the task to reschedule. 918 @return: ID of the newly created task. 919 """ 920 task = self.view_task(task_id) 921 922 if not task: 923 return None 924 925 if task.category == "file": 926 add = self.add_path 927 elif task.category == "url": 928 add = self.add_url 929 930 # Change status to recovered. 931 session = self.Session() 932 session.query(Task).get(task_id).status = TASK_RECOVERED 933 try: 934 session.commit() 935 except SQLAlchemyError as e: 936 log.debug("Database error rescheduling task: {0}".format(e)) 937 session.rollback() 938 return False 939 finally: 940 session.close() 941 942 # Normalize tags. 943 if task.tags: 944 tags = ",".join([tag.name for tag in task.tags]) 945 else: 946 tags = task.tags 947 948 return add(task.target, 949 task.timeout, 950 task.package, 951 task.options, 952 task.priority, 953 task.custom, 954 task.machine, 955 task.platform, 956 tags, 957 task.memory, 958 task.enforce_timeout, 959 task.clock)
960
961 - def list_tasks(self, limit=None, details=False, category=None, offset=None, status=None, not_status=None):
962 """Retrieve list of task. 963 @param limit: specify a limit of entries. 964 @param details: if details about must be included 965 @param category: filter by category 966 @param offset: list offset 967 @param status: filter by task status 968 @param not_status: exclude this task status from filter 969 @return: list of tasks. 970 """ 971 session = self.Session() 972 try: 973 search = session.query(Task) 974 975 if status: 976 search = search.filter(Task.status == status) 977 if not_status: 978 search = search.filter(Task.status != not_status) 979 if category: 980 search = search.filter(Task.category == category) 981 if details: 982 search = search.options(joinedload("guest"), joinedload("errors"), joinedload("tags")) 983 984 tasks = search.order_by("added_on desc").limit(limit).offset(offset).all() 985 except SQLAlchemyError as e: 986 log.debug("Database error listing tasks: {0}".format(e)) 987 return None 988 finally: 989 session.close() 990 return tasks
991
992 - def count_tasks(self, status=None):
993 """Count tasks in the database 994 @param status: apply a filter according to the task status 995 @return: number of tasks found 996 """ 997 session = self.Session() 998 try: 999 if status: 1000 tasks_count = session.query(Task).filter(Task.status == status).count() 1001 else: 1002 tasks_count = session.query(Task).count() 1003 except SQLAlchemyError as e: 1004 log.debug("Database error counting tasks: {0}".format(e)) 1005 return 0 1006 finally: 1007 session.close() 1008 return tasks_count
1009
1010 - def view_task(self, task_id, details=False):
1011 """Retrieve information on a task. 1012 @param task_id: ID of the task to query. 1013 @return: details on the task. 1014 """ 1015 session = self.Session() 1016 try: 1017 if details: 1018 task = session.query(Task).options(joinedload("guest"), joinedload("errors"), joinedload("tags")).get(task_id) 1019 else: 1020 task = session.query(Task).get(task_id) 1021 except SQLAlchemyError as e: 1022 log.debug("Database error viewing task: {0}".format(e)) 1023 return None 1024 else: 1025 if task: 1026 session.expunge(task) 1027 finally: 1028 session.close() 1029 return task
1030
1031 - def delete_task(self, task_id):
1032 """Delete information on a task. 1033 @param task_id: ID of the task to query. 1034 @return: operation status. 1035 """ 1036 session = self.Session() 1037 try: 1038 task = session.query(Task).get(task_id) 1039 session.delete(task) 1040 session.commit() 1041 except SQLAlchemyError as e: 1042 log.debug("Database error deleting task: {0}".format(e)) 1043 session.rollback() 1044 return False 1045 finally: 1046 session.close() 1047 return True
1048
1049 - def view_sample(self, sample_id):
1050 """Retrieve information on a sample given a sample id. 1051 @param sample_id: ID of the sample to query. 1052 @return: details on the sample used in sample: sample_id. 1053 """ 1054 session = self.Session() 1055 try: 1056 sample = session.query(Sample).get(sample_id) 1057 except AttributeError: 1058 return None 1059 except SQLAlchemyError as e: 1060 log.debug("Database error viewing task: {0}".format(e)) 1061 return None 1062 else: 1063 if sample: 1064 session.expunge(sample) 1065 finally: 1066 session.close() 1067 1068 return sample
1069
1070 - def find_sample(self, md5=None, sha256=None):
1071 """Search samples by MD5. 1072 @param md5: md5 string 1073 @return: matches list 1074 """ 1075 session = self.Session() 1076 try: 1077 if md5: 1078 sample = session.query(Sample).filter(Sample.md5 == md5).first() 1079 elif sha256: 1080 sample = session.query(Sample).filter(Sample.sha256 == sha256).first() 1081 except SQLAlchemyError as e: 1082 log.debug("Database error searching sample: {0}".format(e)) 1083 return None 1084 else: 1085 if sample: 1086 session.expunge(sample) 1087 finally: 1088 session.close() 1089 return sample
1090
1091 - def count_samples(self):
1092 """Counts the amount of samples in the database.""" 1093 session = self.Session() 1094 try: 1095 sample_count = session.query(Sample).count() 1096 except SQLAlchemyError as e: 1097 log.debug("Database error counting samples: {0}".format(e)) 1098 return 0 1099 finally: 1100 session.close() 1101 return sample_count
1102
1103 - def view_machine(self, name):
1104 """Show virtual machine. 1105 @params name: virtual machine name 1106 @return: virtual machine's details 1107 """ 1108 session = self.Session() 1109 try: 1110 machine = session.query(Machine).options(joinedload("tags")).filter(Machine.name == name).first() 1111 except SQLAlchemyError as e: 1112 log.debug("Database error viewing machine: {0}".format(e)) 1113 return None 1114 else: 1115 if machine: 1116 session.expunge(machine) 1117 finally: 1118 session.close() 1119 return machine
1120
1121 - def view_machine_by_label(self, label):
1122 """Show virtual machine. 1123 @params label: virtual machine label 1124 @return: virtual machine's details 1125 """ 1126 session = self.Session() 1127 try: 1128 machine = session.query(Machine).options(joinedload("tags")).filter(Machine.label == label).first() 1129 except SQLAlchemyError as e: 1130 log.debug("Database error viewing machine by label: {0}".format(e)) 1131 return None 1132 else: 1133 if machine: 1134 session.expunge(machine) 1135 finally: 1136 session.close() 1137 return machine
1138
1139 - def view_errors(self, task_id):
1140 """Get all errors related to a task. 1141 @param task_id: ID of task associated to the errors 1142 @return: list of errors. 1143 """ 1144 session = self.Session() 1145 try: 1146 errors = session.query(Error).filter(Error.task_id == task_id).all() 1147 except SQLAlchemyError as e: 1148 log.debug("Database error viewing errors: {0}".format(e)) 1149 return None 1150 finally: 1151 session.close() 1152 return errors
1153