示例#1
0
def add_file(file_path, name=None, tags=None, parent=None):
    obj = File(file_path)
    new_path = store_sample(obj, __project__)
    print(new_path)

    if not name:
        name = os.path.basename(file_path)

    # success = True
    if new_path:
        # Add file to the database.
        try:
            db = Database()
            db.add(obj=obj, name=name, tags=tags, parent_sha=parent)
        except Exception as e:
            log.error("Exception while adding sample to DB: {str(e)}")
            # Removing stored file since DB write failed
            remove_sample(new_path)
            return None

        # AutoRun Modules
        if cfg.autorun.enabled:
            autorun_module(obj.sha256)
        # Close the open session to keep the session table clean
        __sessions__.close()
        return obj.sha256
    else:
        log.info("File already exists in database")
        return None
示例#2
0
    def parse_message(self, message_folder):
        db = Database()
        email_header = os.path.join(message_folder, 'InternetHeaders.txt')
        email_body = os.path.join(message_folder, 'Message.txt')

        envelope = headers = email_text = ''
        if os.path.exists(email_header):
            envelope, headers = self.email_headers(email_header)
        if os.path.exists(email_body):
            email_text = open(email_body, 'rb').read()

        tags = 'pst, {0}'.format(message_folder)
        if os.path.exists(os.path.join(message_folder, 'Attachments')):
            for filename in os.listdir(
                    os.path.join(message_folder, 'Attachments')):
                if os.path.isfile(
                        os.path.join(message_folder, 'Attachments', filename)):
                    obj = File(
                        os.path.join(message_folder, 'Attachments', filename))
                    sha256 = hashlib.sha256(
                        open(
                            os.path.join(message_folder, 'Attachments',
                                         filename), 'rb').read()).hexdigest()
                    new_path = store_sample(obj)
                    if new_path:
                        # Add file to the database.
                        db.add(obj=obj, tags=tags)
                    # Add Email Details as a Note
                    # To handle duplicates we use multiple notes
                    headers_body = 'Envelope: \n{0}\nHeaders: \n{1}\n'.format(
                        envelope, headers)
                    db.add_note(sha256, 'Headers', headers_body)

                    # Add a note with email body
                    db.add_note(sha256, 'Email Body', string_clean(email_text))
示例#3
0
文件: pst.py 项目: AnyMaster/viper
 def parse_message(self, message_folder):
     db = Database()
     email_header = os.path.join(message_folder, 'InternetHeaders.txt')
     email_body = os.path.join(message_folder, 'Message.txt')
     
     envelope = headers = email_text = ''
     if os.path.exists(email_header):
         envelope, headers = self.email_headers(email_header)
     if os.path.exists(email_body):
         email_text = open(email_body, 'rb').read()
     
     tags = 'pst, {0}'.format(message_folder)
     if os.path.exists(os.path.join(message_folder, 'Attachments')):
         for filename in os.listdir(os.path.join(message_folder, 'Attachments')):
             if os.path.isfile(os.path.join(message_folder, 'Attachments', filename)):
                 obj = File(os.path.join(message_folder, 'Attachments', filename))
                 sha256 = hashlib.sha256(open(os.path.join(message_folder, 'Attachments', filename), 'rb').read()).hexdigest()
                 new_path = store_sample(obj)
                 if new_path:
                     # Add file to the database.
                     db.add(obj=obj, tags=tags)
                 # Add Email Details as a Note
                 # To handle duplicates we use multiple notes
                 headers_body = 'Envelope: \n{0}\nHeaders: \n{1}\n'.format(envelope, headers)
                 db.add_note(sha256, 'Headers', headers_body)
                 
                 # Add a note with email body
                 db.add_note(sha256, 'Email Body', string_clean(email_text))
示例#4
0
 def add_file(self, file_path, tags, parent):
     obj = File(file_path)
     new_path = store_sample(obj)
     if new_path:
         # Add file to the database.
         db = Database()
         db.add(obj=obj, tags=tags, parent_sha=parent)
         return obj.sha256
示例#5
0
文件: cuckoo.py 项目: dgrif/viper
 def add_file(self, file_path, tags, parent):
     obj = File(file_path)
     new_path = store_sample(obj)
     if new_path:
         # Add file to the database.
         db = Database()
         db.add(obj=obj, tags=tags, parent_sha=parent)
         return obj.sha256
示例#6
0
def _add_file(file_path, name, tags, parent_sha):
    obj = File(file_path)
    new_path = store_sample(obj)
    if new_path:
        db = Database()
        db.add(obj=obj, name=name, tags=tags, parent_sha=parent_sha)
        return obj.sha256
    else:
        return None
示例#7
0
def _add_file(file_path, name, tags, parent_sha):
    obj = File(file_path)
    new_path = store_sample(obj)
    if new_path:
        db = Database()
        db.add(obj=obj, name=name, tags=tags, parent_sha=parent_sha)
        return obj.sha256
    else:
        return None
示例#8
0
文件: web.py 项目: lksnyder0/viper
def url_download():
    url = request.forms.get('url')
    tags = request.forms.get('tag_list')
    tags = "url," + tags
    if request.forms.get("tor"):
        upload = network.download(url, tor=True)
    else:
        upload = network.download(url, tor=False)
    if upload == None:
        return template('error.tpl', error="server can't download from URL")
    # Set Project
    project = 'Main'
    db = Database()
    tf = tempfile.NamedTemporaryFile()
    tf.write(upload)
    if tf == None:
        return template('error.tpl', error="server can't download from URL")
    tf.flush()
    tf_obj = File(tf.name)
    tf_obj.name = tf_obj.sha256
    new_path = store_sample(tf_obj)
    success = False
    if new_path:
        # Add file to the database.
        success = db.add(obj=tf_obj, tags=tags)

    if success:
        #redirect("/project/{0}".format(project))
        redirect("/file/Main/" + tf_obj.sha256)
    else:
        return template('error.tpl',
                        error="Unable to Store The File,already in database")
示例#9
0
文件: web.py 项目: blaquee/viper
def url_download():
    url = request.forms.get('url')
    tags = request.forms.get('tag_list')
    tags = "url,"+tags
    if request.forms.get("tor"):
        upload = network.download(url,tor=True)
    else:
        upload = network.download(url,tor=False)
    if upload == None:
        return template('error.tpl', error="server can't download from URL")
    # Set Project
    project = 'Main'
    db = Database()
    tf = tempfile.NamedTemporaryFile()
    tf.write(upload)
    if tf == None:
        return template('error.tpl', error="server can't download from URL")
    tf.flush()
    tf_obj = File(tf.name)
    tf_obj.name = tf_obj.sha256
    new_path = store_sample(tf_obj)
    success = False
    if new_path:
        # Add file to the database.
        success = db.add(obj=tf_obj, tags=tags)

    if success:
        #redirect("/project/{0}".format(project))
        redirect("/file/Main/"+tf_obj.sha256)
    else:
        return template('error.tpl', error="Unable to Store The File,already in database")
示例#10
0
def add_file():
    tags = request.forms.get('tag_list')
    uploads = request.files.getlist('file')

    # Set Project
    project = request.forms.get('project')
    if project in project_list():
        __project__.open(project)
    else:
        __project__.open('../')
        project = 'Main'
    db = Database()
    file_list = []
    # Write temp file to disk
    with upload_temp() as temp_dir:
        for upload in uploads:
            file_path = os.path.join(temp_dir, upload.filename)
            with open(file_path, 'w') as tmp_file:
                tmp_file.write(upload.file.read())
            # Zip Files
            if request.forms.get('unzip'):
                zip_pass = request.forms.get('zip_pass')
                try:
                    with ZipFile(file_path) as zf:
                        zf.extractall(temp_dir, pwd=zip_pass)
                    for root, dirs, files in os.walk(temp_dir, topdown=False):
                        for name in files:
                            if not name == upload.filename:
                                file_list.append(os.path.join(root, name))
                except Exception as e:
                    return template('error.tpl',
                                    error="Error with zipfile - {0}".format(e))
            # Non zip files
            else:
                file_list.append(file_path)

        # Add each file
        for new_file in file_list:
            print new_file
            obj = File(new_file)
            new_path = store_sample(obj)
            success = True
            if new_path:
                # Add file to the database.
                success = db.add(obj=obj, tags=tags)
                if not success:
                    return template(
                        'error.tpl',
                        error="Unable to Store The File: {0}".format(
                            upload.filename))
    redirect("/project/{0}".format(project))
示例#11
0
def add_file(file_path, name=None, url=None, tags=None, parent=None):
    obj = File(file_path, url)
    new_path = store_sample(obj)
    print(new_path)

    if not name:
        name = os.path.basename(file_path)

    # success = True
    if new_path:
        # Add file to the database.
        db = Database()
        db.add(obj=obj, name=name, tags=tags, url=url, parent_sha=parent)
        # AutoRun Modules
        if cfg.autorun.enabled:
            autorun_module(obj.sha256)
            # Close the open session to keep the session table clean
            __sessions__.close()
        return obj.sha256

    else:
        # ToDo Remove the stored file if we cant write to DB
        return
示例#12
0
文件: views.py 项目: Rafiot/viper
def add_file(file_path, name=None, tags=None, parent=None):
    obj = File(file_path)
    new_path = store_sample(obj)
    print(new_path)

    if not name:
        name = os.path.basename(file_path)

    # success = True
    if new_path:
        # Add file to the database.
        db = Database()
        db.add(obj=obj, name=name, tags=tags, parent_sha=parent)

        # AutoRun Modules
        if cfg.autorun.enabled:
            autorun_module(obj.sha256)
            # Close the open session to keep the session table clean
            __sessions__.close()
        return obj.sha256

    else:
        # ToDo Remove the stored file if we cant write to DB
        return
示例#13
0
文件: web.py 项目: pig123/viper
def add_file():
    tags = request.forms.get('tag_list')
    upload = request.files.get('file')

    # Set Project
    project = request.forms.get('project')
    if project in project_list():
        __project__.open(project)
    else:
        __project__.open('../')
        project = 'Main'
    db = Database()    

    # Write temp file to disk
    with upload_temp() as temp_dir:
        file_path = os.path.join(temp_dir, upload.filename)
        with open(file_path, 'w') as tmp_file:
            tmp_file.write(upload.file.read())
        file_list = []
        # Zip Files
        if request.forms.get('unzip'):
            zip_pass = request.forms.get('zip_pass')
            try:
                with ZipFile(file_path) as zf:
                    zf.extractall(temp_dir, pwd=zip_pass)            
                for root, dirs, files in os.walk(temp_dir, topdown=False):
                    for name in files:
                        if not name == upload.filename:
                            file_list.append(os.path.join(root, name))
            except Exception as e:
                return template('error.tpl', error="Error with zipfile - {0}".format(e))
        # Non zip files
        else:
            file_list.append(file_path)
        
        # Add each file
        for new_file in file_list:
            obj = File(new_file)
            new_path = store_sample(obj)
            success = False
            if new_path:
                # Add file to the database.
                success = db.add(obj=obj, tags=tags)
    if success:
        redirect("/project/{0}".format(project))
    else:
        return template('error.tpl', error="Unable to Store The File")
示例#14
0
class Commands(object):
    output = []

    def __init__(self):
        # Open connection to the database.
        self.db = Database()

        # Map commands to their related functions.
        self.commands = dict(
            help=dict(obj=self.cmd_help, description="Show this help message"),
            open=dict(obj=self.cmd_open, description="Open a file"),
            new=dict(obj=self.cmd_new, description="Create new file"),
            close=dict(obj=self.cmd_close, description="Close the current session"),
            info=dict(obj=self.cmd_info, description="Show information on the opened file"),
            notes=dict(obj=self.cmd_notes, description="View, add and edit notes on the opened file"),
            clear=dict(obj=self.cmd_clear, description="Clear the console"),
            store=dict(obj=self.cmd_store, description="Store the opened file to the local repository"),
            delete=dict(obj=self.cmd_delete, description="Delete the opened file"),
            find=dict(obj=self.cmd_find, description="Find a file"),
            tags=dict(obj=self.cmd_tags, description="Modify tags of the opened file"),
            sessions=dict(obj=self.cmd_sessions, description="List or switch sessions"),
            stats=dict(obj=self.cmd_stats, description="Viper Collection Statistics"),
            projects=dict(obj=self.cmd_projects, description="List or switch existing projects"),
            parent=dict(obj=self.cmd_parent, description="Add or remove a parent file"),
            export=dict(obj=self.cmd_export, description="Export the current session to file or zip"),
            analysis=dict(obj=self.cmd_analysis, description="View the stored analysis"),
            rename=dict(obj=self.cmd_rename, description="Rename the file in the database"),
        )

    # Output Logging
    def log(self, event_type, event_data):
        self.output.append(dict(
            type=event_type,
            data=event_data
        ))
        out.print_output([{'type': event_type, 'data': event_data}], console_output['filename'])


    ##
    # CLEAR
    #
    # This command simply clears the shell.
    def cmd_clear(self, *args):
        os.system('clear')

    ##
    # HELP
    #
    # This command simply prints the help message.
    # It lists both embedded commands and loaded modules.
    def cmd_help(self, *args):
        self.log('info', "Commands")

        rows = []
        for command_name, command_item in self.commands.items():
            rows.append([command_name, command_item['description']])

        rows.append(["exit, quit", "Exit Viper"])
        rows = sorted(rows, key=lambda entry: entry[0])

        self.log('table', dict(header=['Command', 'Description'], rows=rows))
        self.log('info', "Modules")

        rows = []
        for module_name, module_item in __modules__.items():
            rows.append([module_name, module_item['description']])

        rows = sorted(rows, key=lambda entry: entry[0])

        self.log('table', dict(header=['Command', 'Description'], rows=rows))

    ##
    # NEW
    #
    # This command is used to create a new session on a new file,
    # useful for copy & paste of content like Email headers

    def cmd_new(self, *args):
        title = input("Enter a title for the new file: ")
        # Create a new temporary file.
        tmp = tempfile.NamedTemporaryFile(delete=False)
        # Open the temporary file with the default editor, or with nano.
        os.system('"${EDITOR:-nano}" ' + tmp.name)
        __sessions__.new(tmp.name)
        __sessions__.current.file.name = title
        self.log('info', "New file with title \"{0}\" added to the current session".format(bold(title)))

    ##
    # OPEN
    #
    # This command is used to open a session on a given file.
    # It either can be an external file path, or a SHA256 hash of a file which
    # has been previously imported and stored.
    # While the session is active, every operation and module executed will be
    # run against the file specified.
    def cmd_open(self, *args):
        parser = argparse.ArgumentParser(prog='open', description="Open a file", epilog="You can also specify a MD5 or SHA256 hash to a previously stored file in order to open a session on it.")
        group = parser.add_mutually_exclusive_group()
        group.add_argument('-f', '--file', action='store_true', help="Target is a file")
        group.add_argument('-u', '--url', action='store_true', help="Target is a URL")
        group.add_argument('-l', '--last', action='store_true', help="Target is the entry number from the last find command's results")
        parser.add_argument('-t', '--tor', action='store_true', help="Download the file through Tor")
        parser.add_argument("value", metavar='PATH, URL, HASH or ID', nargs='*', help="Target to open. Hash can be md5 or sha256. ID has to be from the last search.")

        try:
            args = parser.parse_args(args)
        except:
            return

        target = " ".join(args.value)

        if not args.last and target is None:
            parser.print_usage()
            return

        # If it's a file path, open a session on it.
        if args.file:
            target = os.path.expanduser(target)

            if not os.path.exists(target) or not os.path.isfile(target):
                self.log('error', "File not found: {0}".format(target))
                return

            __sessions__.new(target)
        # If it's a URL, download it and open a session on the temporary file.
        elif args.url:
            data = download(url=target, tor=args.tor)

            if data:
                tmp = tempfile.NamedTemporaryFile(delete=False)
                tmp.write(data)
                tmp.close()

                __sessions__.new(tmp.name)
        # Try to open the specified file from the list of results from
        # the last find command.
        elif args.last:
            if __sessions__.find:
                count = 1
                for item in __sessions__.find:
                    if count == int(target):
                        __sessions__.new(get_sample_path(item.sha256))
                        break

                    count += 1
            else:
                self.log('warning', "You haven't performed a find yet")
        # Otherwise we assume it's an hash of an previously stored sample.
        else:
            target = target.strip().lower()

            if len(target) == 32:
                key = 'md5'
            elif len(target) == 64:
                key = 'sha256'
            else:
                parser.print_usage()
                return

            rows = self.db.find(key=key, value=target)

            if not rows:
                self.log('warning', "No file found with the given hash {0}".format(target))
                return

            path = get_sample_path(rows[0].sha256)
            if path:
                __sessions__.new(path)

    ##
    # CLOSE
    #
    # This command resets the open session.
    # After that, all handles to the opened file should be closed and the
    # shell should be restored to the default prompt.
    def cmd_close(self, *args):
        __sessions__.close()

    ##
    # INFO
    #
    # This command returns information on the open session. It returns details
    # on the file (e.g. hashes) and other information that might available from
    # the database.
    def cmd_info(self, *args):
        if __sessions__.is_set():
            self.log('table', dict(
                header=['Key', 'Value'],
                rows=[
                    ['Name', __sessions__.current.file.name],
                    ['Tags', __sessions__.current.file.tags],
                    ['Path', __sessions__.current.file.path],
                    ['Size', __sessions__.current.file.size],
                    ['Type', __sessions__.current.file.type],
                    ['Mime', __sessions__.current.file.mime],
                    ['MD5', __sessions__.current.file.md5],
                    ['SHA1', __sessions__.current.file.sha1],
                    ['SHA256', __sessions__.current.file.sha256],
                    ['SHA512', __sessions__.current.file.sha512],
                    ['SSdeep', __sessions__.current.file.ssdeep],
                    ['CRC32', __sessions__.current.file.crc32],
                    ['Parent', __sessions__.current.file.parent],
                    ['Children', __sessions__.current.file.children]
                ]
            ))

    ##
    # NOTES
    #
    # This command allows you to view, add, modify and delete notes associated
    # with the currently opened file.
    def cmd_notes(self, *args):
        parser = argparse.ArgumentParser(prog="notes", description="Show information on the opened file")
        group = parser.add_mutually_exclusive_group()
        group.add_argument('-l', '--list', action='store_true', help="List all notes available for the current file")
        group.add_argument('-a', '--add', action='store_true', help="Add a new note to the current file")
        group.add_argument('-v', '--view', metavar='NOTE ID', type=int, help="View the specified note")
        group.add_argument('-e', '--edit', metavar='NOTE ID', type=int, help="Edit an existing note")
        group.add_argument('-d', '--delete', metavar='NOTE ID', type=int, help="Delete an existing note")

        try:
            args = parser.parse_args(args)
        except:
            return

        if not __sessions__.is_set():
            self.log('error', "No open session")
            return

        # check if the file is already stores, otherwise exit as no notes command will work if the file is not stored in the database
        malware = Database().find(key='sha256', value=__sessions__.current.file.sha256)
        if not malware:
            self.log('error', "The opened file doesn't appear to be in the database, have you stored it yet?")
            return

        if args.list:
            # Retrieve all notes for the currently opened file.

            notes = malware[0].note
            if not notes:
                self.log('info', "No notes available for this file yet")
                return

            # Populate table rows.
            rows = [[note.id, note.title] for note in notes]

            # Display list of existing notes.
            self.log('table', dict(header=['ID', 'Title'], rows=rows))

        elif args.add:
            title = input("Enter a title for the new note: ")

            # Create a new temporary file.
            tmp = tempfile.NamedTemporaryFile(delete=False)
            # Open the temporary file with the default editor, or with nano.
            os.system('"${EDITOR:-nano}" ' + tmp.name)
            # Once the user is done editing, we need to read the content and
            # store it in the database.
            body = tmp.read()
            Database().add_note(__sessions__.current.file.sha256, title, body)
            # Finally, remove the temporary file.
            os.remove(tmp.name)

            self.log('info', "New note with title \"{0}\" added to the current file".format(bold(title)))

        elif args.view:
            # Retrieve note wth the specified ID and print it.
            note = Database().get_note(args.view)
            if note:
                self.log('info', bold('Title: ') + note.title)
                self.log('info', bold('Body:') + '\n' + note.body)
            else:
                self.log('info', "There is no note with ID {0}".format(args.view))

        elif args.edit:
            # Retrieve note with the specified ID.
            note = Database().get_note(args.edit)
            if note:
                # Create a new temporary file.
                tmp = tempfile.NamedTemporaryFile(delete=False)
                # Write the old body to the temporary file.
                tmp.write(note.body)
                tmp.close()
                # Open the old body with the text editor.
                os.system('"${EDITOR:-nano}" ' + tmp.name)
                # Read the new body from the temporary file.
                body = open(tmp.name, 'r').read()
                # Update the note entry with the new body.
                Database().edit_note(args.edit, body)
                # Remove the temporary file.
                os.remove(tmp.name)

                self.log('info', "Updated note with ID {0}".format(args.edit))

        elif args.delete:
            # Delete the note with the specified ID.
            Database().delete_note(args.delete)
        else:
            parser.print_usage()

    ##
    # ANALYSIS
    #
    # This command allows you to view the stored output from modules that have been run
    # with the currently opened file.
    def cmd_analysis(self, *args):
        parser = argparse.ArgumentParser(prog="analysis", description="Show stored module results")
        group = parser.add_mutually_exclusive_group()
        group.add_argument('-l', '--list', action='store_true',
                           help="List all module results available for the current file")
        group.add_argument('-v', '--view', metavar='ANALYSIS ID', type=int, help="View the specified analysis")
        group.add_argument('-d', '--delete', metavar='ANALYSIS ID', type=int, help="Delete an existing analysis")

        try:
            args = parser.parse_args(args)
        except:
            return

        if not __sessions__.is_set():
            self.log('error', "No open session")
            return

        # check if the file is already stores, otherwise exit
        malware = Database().find(key='sha256', value=__sessions__.current.file.sha256)
        if not malware:
            self.log('error', "The opened file doesn't appear to be in the database, have you stored it yet?")
            return

        if args.list:
            # Retrieve all analysis for the currently opened file.

            analysis_list = malware[0].analysis
            if not analysis_list:
                self.log('info', "No analysis available for this file yet")
                return

            # Populate table rows.
            rows = [[analysis.id, analysis.cmd_line, analysis.stored_at] for analysis in analysis_list]

            # Display list of existing results.
            self.log('table', dict(header=['ID', 'Cmd Line', 'Saved On'], rows=rows))

        elif args.view:
            # Retrieve analysis wth the specified ID and print it.
            result = Database().get_analysis(args.view)
            if result:
                self.log('info', bold('Cmd Line: ') + result.cmd_line)
                for line in json.loads(result.results):
                    self.log(line['type'], line['data'])
            else:
                self.log('info', "There is no analysis with ID {0}".format(args.view))

    ##
    # STORE
    #
    # This command stores the opened file in the local repository and tries
    # to store details in the database.
    def cmd_store(self, *args):
        parser = argparse.ArgumentParser(prog='store', description="Store the opened file to the local repository")
        parser.add_argument('-d', '--delete', action='store_true', help="Delete the original file")
        parser.add_argument('-f', '--folder', type=str, nargs='+', help="Specify a folder to import")
        parser.add_argument('-s', '--file-size', type=int, help="Specify a maximum file size")
        parser.add_argument('-y', '--file-type', type=str, help="Specify a file type pattern")
        parser.add_argument('-n', '--file-name', type=str, help="Specify a file name pattern")
        parser.add_argument('-t', '--tags', type=str, nargs='+', help="Specify a list of comma-separated tags")

        try:
            args = parser.parse_args(args)
        except:
            return

        if args.folder is not None:
            # Allows to have spaces in the path.
            args.folder = " ".join(args.folder)

        if args.tags is not None:
            # Remove the spaces in the list of tags
            args.tags = "".join(args.tags)

        def add_file(obj, tags=None):
            if get_sample_path(obj.sha256):
                self.log('warning', "Skip, file \"{0}\" appears to be already stored".format(obj.name))
                return False

            if __sessions__.is_attached_misp(quiet=True):
                if tags is not None:
                    tags += ',misp:{}'.format(__sessions__.current.misp_event.event.id)
                else:
                    tags = 'misp:{}'.format(__sessions__.current.misp_event.event.id)

            # Try to store file object into database.
            status = self.db.add(obj=obj, tags=tags)
            if status:
                # If succeeds, store also in the local repository.
                # If something fails in the database (for example unicode strings)
                # we don't want to have the binary lying in the repository with no
                # associated database record.
                new_path = store_sample(obj)
                self.log("success", "Stored file \"{0}\" to {1}".format(obj.name, new_path))

            else:
                return False

            # Delete the file if requested to do so.
            if args.delete:
                try:
                    os.unlink(obj.path)
                except Exception as e:
                    self.log('warning', "Failed deleting file: {0}".format(e))

            return True

        # If the user specified the --folder flag, we walk recursively and try
        # to add all contained files to the local repository.
        # This is note going to open a new session.
        # TODO: perhaps disable or make recursion optional?
        if args.folder is not None:
            # Check if the specified folder is valid.
            if os.path.isdir(args.folder):
                # Walk through the folder and subfolders.
                for dir_name, dir_names, file_names in walk(args.folder):
                    # Add each collected file.
                    for file_name in file_names:
                        file_path = os.path.join(dir_name, file_name)

                        if not os.path.exists(file_path):
                            continue
                        # Check if file is not zero.
                        if not os.path.getsize(file_path) > 0:
                            continue

                        # Check if the file name matches the provided pattern.
                        if args.file_name:
                            if not fnmatch.fnmatch(file_name, args.file_name):
                                # self.log('warning', "Skip, file \"{0}\" doesn't match the file name pattern".format(file_path))
                                continue

                        # Check if the file type matches the provided pattern.
                        if args.file_type:
                            if args.file_type not in File(file_path).type:
                                # self.log('warning', "Skip, file \"{0}\" doesn't match the file type".format(file_path))
                                continue

                        # Check if file exceeds maximum size limit.
                        if args.file_size:
                            # Obtain file size.
                            if os.path.getsize(file_path) > args.file_size:
                                self.log('warning', "Skip, file \"{0}\" is too big".format(file_path))
                                continue

                        file_obj = File(file_path)

                        # Add file.
                        add_file(file_obj, args.tags)
                        if add_file and cfg.autorun.enabled:
                            autorun_module(file_obj.sha256)
                            # Close the open session to keep the session table clean
                            __sessions__.close()

            else:
                self.log('error', "You specified an invalid folder: {0}".format(args.folder))
        # Otherwise we try to store the currently opened file, if there is any.
        else:
            if __sessions__.is_set():
                if __sessions__.current.file.size == 0:
                    self.log('warning', "Skip, file \"{0}\" appears to be empty".format(__sessions__.current.file.name))
                    return False

                # Add file.
                if add_file(__sessions__.current.file, args.tags):
                    # Open session to the new file.
                    self.cmd_open(*[__sessions__.current.file.sha256])
                    if cfg.autorun.enabled:
                        autorun_module(__sessions__.current.file.sha256)
            else:
                self.log('error', "No open session")

    ##
    # RENAME
    #
    # This command renames the currently opened file in the database.
    def cmd_rename(self, *args):
        if __sessions__.is_set():
            if not __sessions__.current.file.id:
                self.log('error', "The opened file does not have an ID, have you stored it yet?")
                return

            self.log('info', "Current name is: {}".format(bold(__sessions__.current.file.name)))

            new_name = input("New name: ")
            if not new_name:
                self.log('error', "File name can't  be empty!")
                return

            self.db.rename(__sessions__.current.file.id, new_name)

            self.log('info', "Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)
        else:
            self.log('error', "No open session")

    ##
    # DELETE
    #
    # This command deletes the currenlty opened file (only if it's stored in
    # the local repository) and removes the details from the database
    def cmd_delete(self, *args):
        parser = argparse.ArgumentParser(prog='delete', description="Delete a file")
        parser.add_argument('-a', '--all', action='store_true', help="Delete ALL files in this project")
        parser.add_argument('-f', '--find', action="store_true", help="Delete ALL files from last find")

        try:
            args = parser.parse_args(args)
        except:
            return

        while True:
            choice = input("Are you sure? It can't be reverted! [y/n] ")
            if choice == 'y':
                break
            elif choice == 'n':
                return

        if args.all:
            if __sessions__.is_set():
                __sessions__.close()

            samples = self.db.find('all')
            for sample in samples:
                self.db.delete_file(sample.id)
                os.remove(get_sample_path(sample.sha256))

            self.log('info', "Deleted a total of {} files.".format(len(samples)))
        elif args.find:
            if __sessions__.find:
                samples = __sessions__.find
                for sample in samples:
                    self.db.delete_file(sample.id)
                    os.remove(get_sample_path(sample.sha256))
                self.log('info', "Deleted {} files.".format(len(samples)))
            else:
                self.log('error', "No find result")

        else:
            if __sessions__.is_set():
                rows = self.db.find('sha256', __sessions__.current.file.sha256)
                if rows:
                    malware_id = rows[0].id
                    if self.db.delete_file(malware_id):
                        self.log("success", "File deleted")
                    else:
                        self.log('error', "Unable to delete file")

                os.remove(__sessions__.current.file.path)
                __sessions__.close()

                self.log('info', "Deleted opened file.")
            else:
                self.log('error', "No session open, and no --all argument. Nothing to delete.")

    ##
    # FIND
    #
    # This command is used to search for files in the database.
    def cmd_find(self, *args):
        parser = argparse.ArgumentParser(prog='find', description="Find a file")
        group = parser.add_mutually_exclusive_group()
        group.add_argument('-t', '--tags', action='store_true', help="List available tags and quit")
        group.add_argument('type', nargs='?', choices=["all", "latest", "name", "type", "mime", "md5", "sha256", "tag", "note", "any", "ssdeep"], help="Where to search.")
        parser.add_argument("value", nargs='?', help="String to search.")
        try:
            args = parser.parse_args(args)
        except:
            return

        # One of the most useful search terms is by tag. With the --tags
        # argument we first retrieve a list of existing tags and the count
        # of files associated with each of them.
        if args.tags:
            # Retrieve list of tags.
            tags = self.db.list_tags()

            if tags:
                rows = []
                # For each tag, retrieve the count of files associated with it.
                for tag in tags:
                    count = len(self.db.find('tag', tag.tag))
                    rows.append([tag.tag, count])

                # Generate the table with the results.
                header = ['Tag', '# Entries']
                rows.sort(key=lambda x: x[1], reverse=True)
                self.log('table', dict(header=header, rows=rows))
            else:
                self.log('warning', "No tags available")

            return

        # At this point, if there are no search terms specified, return.
        if args.type is None:
            parser.print_usage()
            return

        key = args.type
        if key != 'all' and key != 'latest':
            try:
                # The second argument is the search value.
                value = args.value
            except IndexError:
                self.log('error', "You need to include a search term.")
                return
        else:
            value = None

        # Search all the files matching the given parameters.
        items = self.db.find(key, value)
        if not items:
            return

        # Populate the list of search results.
        rows = []
        count = 1
        for item in items:
            tag = ', '.join([t.tag for t in item.tag if t.tag])
            row = [count, item.name, item.mime, item.md5, tag]
            if key == 'ssdeep':
                row.append(item.ssdeep)
            if key == 'latest':
                row.append(item.created_at)

            rows.append(row)
            count += 1

        # Update find results in current session.
        __sessions__.find = items

        # Generate a table with the results.
        header = ['#', 'Name', 'Mime', 'MD5', 'Tags']
        if key == 'latest':
            header.append('Created At')
        if key == 'ssdeep':
            header.append("Ssdeep")
        self.log("table", dict(header=header, rows=rows))

    ##
    # TAGS
    #
    # This command is used to modify the tags of the opened file.
    def cmd_tags(self, *args):
        parser = argparse.ArgumentParser(prog='tags', description="Modify tags of the opened file")
        parser.add_argument('-a', '--add', metavar='TAG', help="Add tags to the opened file (comma separated)")
        parser.add_argument('-d', '--delete', metavar='TAG', help="Delete a tag from the opened file")
        try:
            args = parser.parse_args(args)
        except:
            return

        # This command requires a session to be opened.
        if not __sessions__.is_set():
            self.log('error', "No open session")
            parser.print_usage()
            return

        # If no arguments are specified, there's not much to do.
        # However, it could make sense to also retrieve a list of existing
        # tags from this command, and not just from the "find" command alone.
        if args.add is None and args.delete is None:
            parser.print_usage()
            return

        # TODO: handle situation where addition or deletion of a tag fail.

        db = Database()
        if not db.find(key='sha256', value=__sessions__.current.file.sha256):
            self.log('error', "The opened file is not stored in the database. "
                "If you want to add it use the `store` command.")
            return

        if args.add:
            # Add specified tags to the database's entry belonging to
            # the opened file.
            db.add_tags(__sessions__.current.file.sha256, args.add)
            self.log('info', "Tags added to the currently opened file")

            # We refresh the opened session to update the attributes.
            # Namely, the list of tags returned by the 'info' command
            # needs to be re-generated, or it wouldn't show the new tags
            # until the existing session is closed a new one is opened.
            self.log('info', "Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

        if args.delete:
            # Delete the tag from the database.
            db.delete_tag(args.delete, __sessions__.current.file.sha256)
            # Refresh the session so that the attributes of the file are
            # updated.
            self.log('info', "Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

    ###
    # SESSION
    #
    # This command is used to list and switch across all the opened sessions.
    def cmd_sessions(self, *args):
        parser = argparse.ArgumentParser(prog='sessions', description="Open a file", epilog="List or switch sessions")
        group = parser.add_mutually_exclusive_group()
        group.add_argument('-l', '--list', action='store_true', help="List all existing sessions")
        group.add_argument('-s', '--switch', type=int, help="Switch to the specified session")

        try:
            args = parser.parse_args(args)
        except:
            return

        if args.list:
            if not __sessions__.sessions:
                self.log('info', "There are no opened sessions")
                return

            rows = []
            for session in __sessions__.sessions:
                current = ''
                if session == __sessions__.current:
                    current = 'Yes'

                rows.append([
                    session.id,
                    session.file.name,
                    session.file.md5,
                    session.created_at,
                    current
                ])

            self.log('info', "Opened Sessions:")
            self.log("table", dict(header=['#', 'Name', 'MD5', 'Created At', 'Current'], rows=rows))
        elif args.switch:
            for session in __sessions__.sessions:
                if args.switch == session.id:
                    __sessions__.switch(session)
                    return

            self.log('warning', "The specified session ID doesn't seem to exist")
        else:
            parser.print_usage()

    ##
    # PROJECTS
    #
    # This command retrieves a list of all projects.
    # You can also switch to a different project.
    def cmd_projects(self, *args):
        parser = argparse.ArgumentParser(prog='projects', description="Open a file", epilog="List or switch existing projects")
        group = parser.add_mutually_exclusive_group()
        group.add_argument('-l', '--list', action='store_true', help="List all existing projects")
        group.add_argument('-s', '--switch', metavar='PROJECT NAME', help="Switch to the specified project")

        try:
            args = parser.parse_args(args)
        except:
            return

        projects_path = os.path.join(expanduser("~"), '.viper', 'projects')

        if not os.path.exists(projects_path):
            self.log('info', "The projects directory does not exist yet")
            return

        if args.list:
            self.log('info', "Projects Available:")

            rows = []
            for project in os.listdir(projects_path):
                project_path = os.path.join(projects_path, project)
                if os.path.isdir(project_path):
                    current = ''
                    if __project__.name and project == __project__.name:
                        current = 'Yes'
                    rows.append([project, time.ctime(os.path.getctime(project_path)), current])

            self.log('table', dict(header=['Project Name', 'Creation Time', 'Current'], rows=rows))
        elif args.switch:
            if __sessions__.is_set():
                __sessions__.close()
                self.log('info', "Closed opened session")

            __project__.open(args.switch)
            self.log('info', "Switched to project {0}".format(bold(args.switch)))

            # Need to re-initialize the Database to open the new SQLite file.
            self.db = Database()
        else:
            self.log('info', parser.print_usage())

    ##
    # EXPORT
    #
    # This command will export the current session to file or zip.
    def cmd_export(self, *args):
        parser = argparse.ArgumentParser(prog='export', description="Export the current session to file or zip")
        parser.add_argument('-z', '--zip', action='store_true', help="Export session in a zip archive")
        parser.add_argument('value', help="path or archive name")

        try:
            args = parser.parse_args(args)
        except:
            return

        # This command requires a session to be opened.
        if not __sessions__.is_set():
            self.log('error', "No open session")
            parser.print_usage()
            return

        # Check for valid export path.
        if args.value is None:
            parser.print_usage()
            return

        # TODO: having for one a folder and for the other a full
        # target path can be confusing. We should perhaps standardize this.

        # Abort if the specified path already exists.
        if os.path.isfile(args.value):
            self.log('error', "File at path \"{0}\" already exists, abort".format(args.value))
            return

        # If the argument chosed so, archive the file when exporting it.
        # TODO: perhaps add an option to use a password for the archive
        # and default it to "infected".
        if args.zip:
            try:
                with ZipFile(args.value, 'w') as export_zip:
                    export_zip.write(__sessions__.current.file.path, arcname=__sessions__.current.file.name)
            except IOError as e:
                self.log('error', "Unable to export file: {0}".format(e))
            else:
                self.log('info', "File archived and exported to {0}".format(args.value))
        # Otherwise just dump it to the given directory.
        else:
            # XXX: Export file with the original file name.
            store_path = os.path.join(args.value, __sessions__.current.file.name)

            try:
                shutil.copyfile(__sessions__.current.file.path, store_path)
            except IOError as e:
                self.log('error', "Unable to export file: {0}".format(e))
            else:
                self.log('info', "File exported to {0}".format(store_path))

    ##
    # STATS
    #
    # This command allows you to generate basic statistics for the stored files.
    def cmd_stats(self, *args):
        parser = argparse.ArgumentParser(prog='stats', description="Display Database File Statistics")
        parser.add_argument('-t', '--top', type=int, help='Top x Items')

        try:
            args = parser.parse_args(args)
        except:
            return

        arg_top = args.top

        # Set all Counters Dict
        extension_dict = defaultdict(int)
        mime_dict = defaultdict(int)
        tags_dict = defaultdict(int)
        size_list = []

        # Find all
        items = self.db.find('all')

        if len(items) < 1:
            self.log('info', "No items in database to generate stats")
            return

        # Sort in to stats
        for item in items:
            if '.' in item.name:
                ext = item.name.split('.')
                extension_dict[ext[-1]] += 1
            mime_dict[item.mime] += 1
            size_list.append(item.size)
            for t in item.tag:
                if t.tag:
                    tags_dict[t.tag] += 1

        avg_size = sum(size_list) / len(size_list)
        #all_stats = {'Total': len(items), 'File Extension': extension_dict, 'Mime': mime_dict, 'Tags': tags_dict,
        #             'Avg Size': avg_size, 'Largest': max(size_list), 'Smallest': min(size_list)}

        # Counter for top x
        if arg_top:
            counter = arg_top
            prefix = 'Top {0} '.format(counter)
        else:
            counter = len(items)
            prefix = ''

        # Project Stats Last as i have it iterate them all

        # Print all the results

        self.log('info', "Projects")
        self.log('table', dict(header=['Name', 'Count'], rows=[['Main', len(items)], ['Next', '10']]))

        # For Current Project
        self.log('info', "Current Project")

        # Extension
        self.log('info', "{0}Extensions".format(prefix))
        header = ['Ext', 'Count']
        rows = []

        for k in sorted(extension_dict, key=extension_dict.get, reverse=True)[:counter]:
            rows.append([k, extension_dict[k]])
        self.log('table', dict(header=header, rows=rows))


        # Mimes
        self.log('info', "{0}Mime Types".format(prefix))
        header = ['Mime', 'Count']
        rows = []
        for k in sorted(mime_dict, key=mime_dict.get, reverse=True)[:counter]:
            rows.append([k, mime_dict[k]])
        self.log('table', dict(header=header, rows=rows))

        # Tags
        self.log('info', "{0}Tags".format(prefix))
        header = ['Tag', 'Count']
        rows = []
        for k in sorted(tags_dict, key=tags_dict.get, reverse=True)[:counter]:
            rows.append([k, tags_dict[k]])
        self.log('table', dict(header=header, rows=rows))

        # Size
        self.log('info', "Size Stats")
        self.log('item', "Largest  {0}".format(convert_size(max(size_list))))
        self.log('item', "Smallest  {0}".format(convert_size(min(size_list))))
        self.log('item', "Average  {0}".format(convert_size(avg_size)))

    ##
    # PARENT
    #
    # This command is used to view or edit the parent child relationship between files.
    def cmd_parent(self, *args):
        parser = argparse.ArgumentParser(prog='tags', description="Set the Parent for this file.")
        parser.add_argument('-a', '--add', metavar='SHA256', help="Add parent file by sha256")
        parser.add_argument('-d', '--delete', action='store_true', help="Delete Parent")
        parser.add_argument('-o', '--open', action='store_true', help="Open The Parent")
        try:
            args = parser.parse_args(args)
        except:
            return

        # This command requires a session to be opened.
        if not __sessions__.is_set():
            self.log('error', "No open session")
            parser.print_usage()
            return


        # If no arguments are specified, there's not much to do.
        if args.add is None and args.delete is None and args.open is None:
            parser.print_usage()
            return

        db = Database()
        if not db.find(key='sha256', value=__sessions__.current.file.sha256):
            self.log('error', "The opened file is not stored in the database. "
                              "If you want to add it use the `store` command.")
            return

        if args.add:
            if not db.find(key='sha256', value=args.add):
                self.log('error', "the parent file is not found in the database. ")
                return
            db.add_parent(__sessions__.current.file.sha256, args.add)
            self.log('info', "parent added to the currently opened file")

            self.log('info', "Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

        if args.delete:
            db.delete_parent(__sessions__.current.file.sha256)
            self.log('info', "parent removed from the currently opened file")

            self.log('info', "Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

        if args.open:
            # Open a session on the parent
            if __sessions__.current.file.parent:
                __sessions__.new(get_sample_path(__sessions__.current.file.parent[-64:]))
            else:
                self.log('info', "No parent set for this sample")
示例#15
0
    def test_add_unicode_py3(self, capsys, filename, name):
        f = File(os.path.join(FIXTURE_DIR, filename))

        instance = Database()
        ret = instance.add(f)
        assert ret is True
示例#16
0
class Commands(object):

    output = []
    
    def __init__(self):
        # Open connection to the database.
        self.db = Database()

        # Map commands to their related functions.
        self.commands = dict(
            help=dict(obj=self.cmd_help, description="Show this help message"),
            open=dict(obj=self.cmd_open, description="Open a file"),
            new=dict(obj=self.cmd_new, description="Create new file"),
            close=dict(obj=self.cmd_close, description="Close the current session"),
            info=dict(obj=self.cmd_info, description="Show information on the opened file"),
            notes=dict(obj=self.cmd_notes, description="View, add and edit notes on the opened file"),
            clear=dict(obj=self.cmd_clear, description="Clear the console"),
            store=dict(obj=self.cmd_store, description="Store the opened file to the local repository"),
            delete=dict(obj=self.cmd_delete, description="Delete the opened file"),
            find=dict(obj=self.cmd_find, description="Find a file"),
            tags=dict(obj=self.cmd_tags, description="Modify tags of the opened file"),
            sessions=dict(obj=self.cmd_sessions, description="List or switch sessions"),
            projects=dict(obj=self.cmd_projects, description="List or switch existing projects"),
            export=dict(obj=self.cmd_export, description="Export the current session to file or zip"),
        )
        
    # Output Logging
    def log(self, event_type, event_data):
        self.output.append(dict(
            type=event_type,
            data=event_data
        ))
        
    ##
    # CLEAR
    #
    # This command simply clears the shell.
    def cmd_clear(self, *args):
        os.system('clear')

    ##
    # HELP
    #
    # This command simply prints the help message.
    # It lists both embedded commands and loaded modules.
    def cmd_help(self, *args):
        self.log('info', "Commands")

        rows = []
        for command_name, command_item in self.commands.items():
            rows.append([command_name, command_item['description']])

        rows.append(["exit, quit", "Exit Viper"])
        rows = sorted(rows, key=lambda entry: entry[0])

        self.log('table', dict(header=['Command', 'Description'], rows=rows))
        self.log('info', "Modules")

        rows = []
        for module_name, module_item in __modules__.items():
            rows.append([module_name, module_item['description']])

        rows = sorted(rows, key=lambda entry: entry[0])

        self.log('table', dict(header=['Command', 'Description'], rows=rows))

    ##
    # NEW
    #
    # This command is used to create a new session on a new file,
    # useful for copy & paste of content like Email headers

    def cmd_new(self, *args):
        title = raw_input("Enter a title for the new file: ")
        # Create a new temporary file.
        tmp = tempfile.NamedTemporaryFile(delete=False)
        # Open the temporary file with the default editor, or with nano.
        os.system('"${EDITOR:-nano}" ' + tmp.name)
        __sessions__.new(tmp.name)
        __sessions__.current.file.name = title
        print_info("New file with title \"{0}\" added to the current session".format(bold(title)))

    ##
    # OPEN
    #
    # This command is used to open a session on a given file.
    # It either can be an external file path, or a SHA256 hash of a file which
    # has been previously imported and stored.
    # While the session is active, every operation and module executed will be
    # run against the file specified.
    def cmd_open(self, *args):
        parser = argparse.ArgumentParser(prog="open", description="Open a file", epilog="You can also specify a MD5 or SHA256 hash to a previously stored file in order to open a session on it.")
        group = parser.add_mutually_exclusive_group()
        group.add_argument('-f', '--file', action="store_true", help="target is a file")
        group.add_argument('-u', '--url', action="store_true", help="target is a URL")
        group.add_argument('-l', '--last', action="store_true", help="target is the entry number from the last find command's results")
        parser.add_argument('-t', '--tor', action="store_true", help="Download the file through Tor")
        parser.add_argument("value", metavar='Path, URL, hash or ID', nargs='*', help="Target to open. Hash can be md5 or sha256. ID has to be from the last search.")

        try:
            args = parser.parse_args(args)
        except:
            return

        target = " ".join(args.value)

        if not args.last and target is None:
            parser.print_usage()
            return

        # If it's a file path, open a session on it.
        if args.file:
            target = os.path.expanduser(target)

            if not os.path.exists(target) or not os.path.isfile(target):
                self.log('error', "File not found: {0}".format(target))
                return

            __sessions__.new(target)
        # If it's a URL, download it and open a session on the temporary file.
        elif args.url:
            data = download(url=target, tor=args.tor)

            if data:
                tmp = tempfile.NamedTemporaryFile(delete=False)
                tmp.write(data)
                tmp.close()

                __sessions__.new(tmp.name)
        # Try to open the specified file from the list of results from
        # the last find command.
        elif args.last:
            if __sessions__.find:
                count = 1
                for item in __sessions__.find:
                    if count == int(target):
                        __sessions__.new(get_sample_path(item.sha256))
                        break

                    count += 1
            else:
                self.log('warning', "You haven't performed a find yet")
        # Otherwise we assume it's an hash of an previously stored sample.
        else:
            target = target.strip().lower()

            if len(target) == 32:
                key = 'md5'
            elif len(target) == 64:
                key = 'sha256'
            else:
                parser.print_usage()
                return

            rows = self.db.find(key=key, value=target)

            if not rows:
                self.log('warning', "No file found with the given hash {0}".format(target))
                return

            path = get_sample_path(rows[0].sha256)
            if path:
                __sessions__.new(path)

    ##
    # CLOSE
    #
    # This command resets the open session.
    # After that, all handles to the opened file should be closed and the
    # shell should be restored to the default prompt.
    def cmd_close(self, *args):
        __sessions__.close()

    ##
    # INFO
    #
    # This command returns information on the open session. It returns details
    # on the file (e.g. hashes) and other information that might available from
    # the database.
    def cmd_info(self, *args):
        if __sessions__.is_set():
            self.log('table', dict(
                header=['Key', 'Value'],
                rows=[
                    ['Name', __sessions__.current.file.name],
                    ['Tags', __sessions__.current.file.tags],
                    ['Path', __sessions__.current.file.path],
                    ['Size', __sessions__.current.file.size],
                    ['Type', __sessions__.current.file.type],
                    ['Mime', __sessions__.current.file.mime],
                    ['MD5', __sessions__.current.file.md5],
                    ['SHA1', __sessions__.current.file.sha1],
                    ['SHA256', __sessions__.current.file.sha256],
                    ['SHA512', __sessions__.current.file.sha512],
                    ['SSdeep', __sessions__.current.file.ssdeep],
                    ['CRC32', __sessions__.current.file.crc32]
                ]
            ))

    ##
    # NOTES
    #
    # This command allows you to view, add, modify and delete notes associated
    # with the currently opened file.
    def cmd_notes(self, *args):
        parser = argparse.ArgumentParser(prog="notes", description="Show information on the opened file")
        group = parser.add_mutually_exclusive_group()
        group.add_argument('-l', '--list', action="store_true", help="List all notes available for the current file")
        group.add_argument('-a', '--add', action="store_true", help="Add a new note to the current file")
        group.add_argument('-v', '--view', metavar='note_id', type=int, help="View the specified note")
        group.add_argument('-e', '--edit', metavar='note_id', type=int, help="Edit an existing note")
        group.add_argument('-d', '--delete', metavar='note_id', type=int, help="Delete an existing note")

        try:
            args = parser.parse_args(args)
        except:
            return

        if not __sessions__.is_set():
            self.log('error', "No session opened")
            return

        if args.list:
            # Retrieve all notes for the currently opened file.
            malware = Database().find(key='sha256', value=__sessions__.current.file.sha256)
            if not malware:
                self.log('error', "The opened file doesn't appear to be in the database, have you stored it yet?")
                return

            notes = malware[0].note
            if not notes:
                self.log('info', "No notes available for this file yet")
                return

            # Populate table rows.
            rows = [[note.id, note.title] for note in notes]

            # Display list of existing notes.
            self.log('table', dict(header=['ID', 'Title'], rows=rows))

        elif args.add:
            title = raw_input("Enter a title for the new note: ")

            # Create a new temporary file.
            tmp = tempfile.NamedTemporaryFile(delete=False)
            # Open the temporary file with the default editor, or with nano.
            os.system('"${EDITOR:-nano}" ' + tmp.name)
            # Once the user is done editing, we need to read the content and
            # store it in the database.
            body = tmp.read()
            Database().add_note(__sessions__.current.file.sha256, title, body)
            # Finally, remove the temporary file.
            os.remove(tmp.name)

            self.log('info', "New note with title \"{0}\" added to the current file".format(bold(title)))

        elif args.view:
            # Retrieve note wth the specified ID and print it.
            note = Database().get_note(args.view)
            if note:
                self.log('info', bold('Title: ') + note.title)
                self.log('info', bold('Body:') + '\n' + note.body)
            else:
                self.log('info', "There is no note with ID {0}".format(args.view))

        elif args.edit:
            # Retrieve note with the specified ID.
            note = Database().get_note(args.edit)
            if note:
                # Create a new temporary file.
                tmp = tempfile.NamedTemporaryFile(delete=False)
                # Write the old body to the temporary file.
                tmp.write(note.body)
                tmp.close()
                # Open the old body with the text editor.
                os.system('"${EDITOR:-nano}" ' + tmp.name)
                # Read the new body from the temporary file.
                body = open(tmp.name, 'r').read()
                # Update the note entry with the new body.
                Database().edit_note(args.edit, body)
                # Remove the temporary file.
                os.remove(tmp.name)

                self.log('info', "Updated note with ID {0}".format(args.edit))

        elif args.delete:
            # Delete the note with the specified ID.
            Database().delete_note(args.delete)
        else:
            parser.print_usage()

    ##
    # STORE
    #
    # This command stores the opened file in the local repository and tries
    # to store details in the database.
    def cmd_store(self, *args):
        parser = argparse.ArgumentParser(prog="store", description="Store the opened file to the local repository")
        parser.add_argument('-d', '--delete', action="store_true", help="Delete the original file")
        parser.add_argument('-f', '--folder', type=str, nargs='+', help="Specify a folder to import")
        parser.add_argument('-s', '--file-size', type=int, help="Specify a maximum file size")
        parser.add_argument('-y', '--file-type', type=str, help="Specify a file type pattern")
        parser.add_argument('-n', '--file-name', type=str, help="Specify a file name pattern")
        parser.add_argument('-t', '--tags', type=str, nargs='+', help="Specify a list of comma-separated tags")

        try:
            args = parser.parse_args(args)
        except:
            return

        if args.folder is not None:
            # Allows to have spaces in the path.
            args.folder = " ".join(args.folder)

        if args.tags is not None:
            # Remove the spaces in the list of tags
            args.tags = "".join(args.tags)

        def add_file(obj, tags=None):
            if get_sample_path(obj.sha256):
                self.log('warning', "Skip, file \"{0}\" appears to be already stored".format(obj.name))
                return False

            # Try to store file object into database.
            status = self.db.add(obj=obj, tags=tags)
            if status:
                # If succeeds, store also in the local repository.
                # If something fails in the database (for example unicode strings)
                # we don't want to have the binary lying in the repository with no
                # associated database record.
                new_path = store_sample(obj)
                self.log("success", "Stored file \"{0}\" to {1}".format(obj.name, new_path))
            else:
                return False

            # Delete the file if requested to do so.
            if args.delete:
                try:
                    os.unlink(obj.path)
                except Exception as e:
                    self.log('warning', "Failed deleting file: {0}".format(e))

            return True

        # If the user specified the --folder flag, we walk recursively and try
        # to add all contained files to the local repository.
        # This is note going to open a new session.
        # TODO: perhaps disable or make recursion optional?
        if args.folder is not None:
            # Check if the specified folder is valid.
            if os.path.isdir(args.folder):
                # Walk through the folder and subfolders.
                for dir_name, dir_names, file_names in os.walk(args.folder):
                    # Add each collected file.
                    for file_name in file_names:
                        file_path = os.path.join(dir_name, file_name)

                        if not os.path.exists(file_path):
                            continue
                        # Check if file is not zero.
                        if not os.path.getsize(file_path) > 0:
                            continue

                        # Check if the file name matches the provided pattern.
                        if args.file_name:
                            if not fnmatch.fnmatch(file_name, args.file_name):
                                # self.log('warning', "Skip, file \"{0}\" doesn't match the file name pattern".format(file_path))
                                continue

                        # Check if the file type matches the provided pattern.
                        if args.file_type:
                            if args.file_type not in File(file_path).type:
                                # self.log('warning', "Skip, file \"{0}\" doesn't match the file type".format(file_path))
                                continue

                        # Check if file exceeds maximum size limit.
                        if args.file_size:
                            # Obtain file size.
                            if os.path.getsize(file_path) > args.file_size:
                                self.log('warning', "Skip, file \"{0}\" is too big".format(file_path))
                                continue

                        file_obj = File(file_path)

                        # Add file.
                        add_file(file_obj, args.tags)
            else:
                self.log('error', "You specified an invalid folder: {0}".format(args.folder))
        # Otherwise we try to store the currently opened file, if there is any.
        else:
            if __sessions__.is_set():
                if __sessions__.current.file.size == 0:
                    self.log('warning', "Skip, file \"{0}\" appears to be empty".format(__sessions__.current.file.name))
                    return False

                # Add file.
                if add_file(__sessions__.current.file, args.tags):
                    # Open session to the new file.
                    self.cmd_open(*[__sessions__.current.file.sha256])
            else:
                self.log('error', "No session opened")

    ##
    # DELETE
    #
    # This commands deletes the currenlty opened file (only if it's stored in
    # the local repository) and removes the details from the database
    def cmd_delete(self, *args):
        if __sessions__.is_set():
            while True:
                choice = raw_input("Are you sure you want to delete this binary? Can't be reverted! [y/n] ")
                if choice == 'y':
                    break
                elif choice == 'n':
                    return

            rows = self.db.find('sha256', __sessions__.current.file.sha256)
            if rows:
                malware_id = rows[0].id
                if self.db.delete_file(malware_id):
                    self.log("success", "File deleted")
                else:
                    self.log('error', "Unable to delete file")

            os.remove(__sessions__.current.file.path)
            __sessions__.close()
        else:
            self.log('error', "No session opened")

    ##
    # FIND
    #
    # This command is used to search for files in the database.
    def cmd_find(self, *args):
        parser = argparse.ArgumentParser(prog="find", description="Find a file")
        group = parser.add_mutually_exclusive_group()
        group.add_argument('-t', '--tags', action="store_true", help="List available tags and quit")
        group.add_argument('type', nargs='?', choices=["all", "latest", "name", "type", "mime", "md5", "sha256", "tag", "note"], help="Where to search.")
        parser.add_argument("value", nargs='?', help="String to search.")
        try:
            args = parser.parse_args(args)
        except:
            return

        # One of the most useful search terms is by tag. With the --tags
        # argument we first retrieve a list of existing tags and the count
        # of files associated with each of them.
        if args.tags:
            # Retrieve list of tags.
            tags = self.db.list_tags()

            if tags:
                rows = []
                # For each tag, retrieve the count of files associated with it.
                for tag in tags:
                    count = len(self.db.find('tag', tag.tag))
                    rows.append([tag.tag, count])

                # Generate the table with the results.
                header = ['Tag', '# Entries']
                rows.sort(key=lambda x: x[1], reverse=True)
                self.log('table', dict(header=header, rows=rows))
            else:
                self.log('warning', "No tags available")

            return

        # At this point, if there are no search terms specified, return.
        if args.type is None:
            parser.print_usage()
            return

        key = args.type
        if key != 'all' and key != 'latest':
            try:
                # The second argument is the search value.
                value = args.value
            except IndexError:
                self.log('error', "You need to include a search term.")
                return
        else:
            value = None

        # Search all the files matching the given parameters.
        items = self.db.find(key, value)
        if not items:
            return

        # Populate the list of search results.
        rows = []
        count = 1
        for item in items:
            tag = ', '.join([t.tag for t in item.tag if t.tag])
            row = [count, item.name, item.mime, item.md5, tag]
            if key == 'latest':
                row.append(item.created_at)

            rows.append(row)
            count += 1

        # Update find results in current session.
        __sessions__.find = items

        # Generate a table with the results.
        header = ['#', 'Name', 'Mime', 'MD5', 'Tags']
        if key == 'latest':
            header.append('Created At')

        self.log("table", dict(header=header, rows=rows))

    ##
    # TAGS
    #
    # This command is used to modify the tags of the opened file.
    def cmd_tags(self, *args):
        parser = argparse.ArgumentParser(prog="tags", description="Modify tags of the opened file")
        parser.add_argument('-a', '--add', help="Add tags to the opened file (comma separated)")
        parser.add_argument('-d', '--delete', help="Delete a tag from the opened file")
        try:
            args = parser.parse_args(args)
        except:
            return

        # This command requires a session to be opened.
        if not __sessions__.is_set():
            self.log('error', "No session opened")
            parser.print_usage()
            return

        # If no arguments are specified, there's not much to do.
        # However, it could make sense to also retrieve a list of existing
        # tags from this command, and not just from the "find" command alone.
        if args.add is None and args.delete is None:
            parser.print_usage()
            return

        # TODO: handle situation where addition or deletion of a tag fail.

        if args.add:
            # Add specified tags to the database's entry belonging to
            # the opened file.
            db = Database()
            db.add_tags(__sessions__.current.file.sha256, args.add)
            self.log('info', "Tags added to the currently opened file")

            # We refresh the opened session to update the attributes.
            # Namely, the list of tags returned by the 'info' command
            # needs to be re-generated, or it wouldn't show the new tags
            # until the existing session is closed a new one is opened.
            self.log('info', "Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

        if args.delete:
            # Delete the tag from the database.
            Database().delete_tag(args.delete)
            # Refresh the session so that the attributes of the file are
            # updated.
            self.log('info', "Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

    ###
    # SESSION
    #
    # This command is used to list and switch across all the opened sessions.
    def cmd_sessions(self, *args):
        parser = argparse.ArgumentParser(prog="sessions", description="Open a file", epilog="List or switch sessions")
        group = parser.add_mutually_exclusive_group()
        group.add_argument('-l', '--list', action="store_true", help="List all existing sessions")
        group.add_argument('-s', '--switch', type=int, help="Switch to the specified session")

        try:
            args = parser.parse_args(args)
        except:
            return

        if args.list:
            if not __sessions__.sessions:
                self.log('info', "There are no opened sessions")
                return

            rows = []
            for session in __sessions__.sessions:
                current = ''
                if session == __sessions__.current:
                    current = 'Yes'

                rows.append([
                    session.id,
                    session.file.name,
                    session.file.md5,
                    session.created_at,
                    current
                ])

            self.log('info', "Opened Sessions:")
            self.log("table", dict(header=['#', 'Name', 'MD5', 'Created At', 'Current'], rows=rows))
        elif args.switch:
            for session in __sessions__.sessions:
                if args.switch == session.id:
                    __sessions__.switch(session)
                    return

            self.log('warning', "The specified session ID doesn't seem to exist")
        else:
            parser.print_usage()

    ##
    # PROJECTS
    #
    # This command retrieves a list of all projects.
    # You can also switch to a different project.
    def cmd_projects(self, *args):
        parser = argparse.ArgumentParser(prog="projects", description="Open a file", epilog="List or switch existing projects")
        group = parser.add_mutually_exclusive_group()
        group.add_argument('-l', '--list', action="store_true", help="List all existing projects")
        group.add_argument('-s', '--switch', metavar='project_name', help="Switch to the specified project")

        try:
            args = parser.parse_args(args)
        except:
            return

        projects_path = os.path.join(os.getcwd(), 'projects')

        if not os.path.exists(projects_path):
            self.log('info', "The projects directory does not exist yet")
            return

        if args.list:
            self.log('info', "Projects Available:")

            rows = []
            for project in os.listdir(projects_path):
                project_path = os.path.join(projects_path, project)
                if os.path.isdir(project_path):
                    current = ''
                    if __project__.name and project == __project__.name:
                        current = 'Yes'
                    rows.append([project, time.ctime(os.path.getctime(project_path)), current])

            self.log('table', dict(header=['Project Name', 'Creation Time', 'Current'], rows=rows))
        elif args.switch:
            if __sessions__.is_set():
                __sessions__.close()
                self.log('info', "Closed opened session")

            __project__.open(args.switch)
            self.log('info', "Switched to project {0}".format(bold(args.switch)))

            # Need to re-initialize the Database to open the new SQLite file.
            self.db = Database()
        else:
            self.log('info', parser.print_usage())

    ##
    # EXPORT
    #
    # This command will export the current session to file or zip.
    def cmd_export(self, *args):
        parser = argparse.ArgumentParser(prog="export", description="Export the current session to file or zip")
        parser.add_argument('-z', '--zip', action="store_true", help="Export session in a zip archive")
        parser.add_argument('value', help="path or archive name")

        try:
            args = parser.parse_args(args)
        except:
            return

        # This command requires a session to be opened.
        if not __sessions__.is_set():
            self.log('error', "No session opened")
            parser.print_usage()
            return

        # Check for valid export path.
        if args.value is None:
            parser.print_usage()
            return

        # TODO: having for one a folder and for the other a full
        # target path can be confusing. We should perhaps standardize this.

        # Abort if the specified path already exists.
        if os.path.isfile(args.value):
            self.log('error', "File at path \"{0}\" already exists, abort".format(args.value))
            return

        # If the argument chosed so, archive the file when exporting it.
        # TODO: perhaps add an option to use a password for the archive
        # and default it to "infected".
        if args.zip:
            try:
                with ZipFile(args.value, 'w') as export_zip:
                    export_zip.write(__sessions__.current.file.path, arcname=__sessions__.current.file.name)
            except IOError as e:
                self.log('error', "Unable to export file: {0}".format(e))
            else:
                self.log('info', "File archived and exported to {0}".format(args.value))
        # Otherwise just dump it to the given directory.
        else:
            # XXX: Export file with the original file name.
            store_path = os.path.join(args.value, __sessions__.current.file.name)

            try:
                shutil.copyfile(__sessions__.current.file.path, store_path)
            except IOError as e:
                self.log('error', "Unable to export file: {0}".format(e))
            else:
                self.log('info', "File exported to {0}".format(store_path))
示例#17
0
class Commands(object):
    def __init__(self):
        # Open connection to the database.
        self.db = Database()

        # Map commands to their related functions.
        self.commands = dict(
            help=dict(obj=self.cmd_help, description="Show this help message"),
            open=dict(obj=self.cmd_open, description="Open a file"),
            close=dict(obj=self.cmd_close, description="Close the current session"),
            info=dict(obj=self.cmd_info, description="Show information on the opened file"),
            notes=dict(obj=self.cmd_notes, description="View, add and edit notes on the opened file"),
            clear=dict(obj=self.cmd_clear, description="Clear the console"),
            store=dict(obj=self.cmd_store, description="Store the opened file to the local repository"),
            delete=dict(obj=self.cmd_delete, description="Delete the opened file"),
            find=dict(obj=self.cmd_find, description="Find a file"),
            tags=dict(obj=self.cmd_tags, description="Modify tags of the opened file"),
            sessions=dict(obj=self.cmd_sessions, description="List or switch sessions"),
            projects=dict(obj=self.cmd_projects, description="List or switch existing projects"),
        )

    ##
    # CLEAR
    #
    # This command simply clears the shell.
    def cmd_clear(self, *args):
        os.system("clear")

    ##
    # HELP
    #
    # This command simply prints the help message.
    # It lists both embedded commands and loaded modules.
    def cmd_help(self, *args):
        print(bold("Commands:"))

        rows = []
        for command_name, command_item in self.commands.items():
            rows.append([command_name, command_item["description"]])

        rows = sorted(rows, key=lambda entry: entry[0])

        print(table(["Command", "Description"], rows))
        print("")
        print(bold("Modules:"))

        rows = []
        for module_name, module_item in __modules__.items():
            rows.append([module_name, module_item["description"]])

        rows = sorted(rows, key=lambda entry: entry[0])

        print(table(["Command", "Description"], rows))

    ##
    # OPEN
    #
    # This command is used to open a session on a given file.
    # It either can be an external file path, or a SHA256 hash of a file which
    # has been previously imported and stored.
    # While the session is active, every operation and module executed will be
    # run against the file specified.
    def cmd_open(self, *args):
        def usage():
            print("usage: open [-h] [-f] [-u] [-l] [-t] <target|md5|sha256>")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--file (-f)\tThe target is a file")
            print("\t--url (-u)\tThe target is a URL")
            print("\t--last (-l)\tOpen file from the results of the last find command")
            print("\t--tor (-t)\tDownload the file through Tor")
            print("")
            print("You can also specify a MD5 or SHA256 hash to a previously stored")
            print("file in order to open a session on it.")
            print("")

        try:
            opts, argv = getopt.getopt(args, "hfult", ["help", "file", "url", "last", "tor"])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_is_file = False
        arg_is_url = False
        arg_last = False
        arg_use_tor = False

        for opt, value in opts:
            if opt in ("-h", "--help"):
                help()
                return
            elif opt in ("-f", "--file"):
                arg_is_file = True
            elif opt in ("-u", "--url"):
                arg_is_url = True
            elif opt in ("-l", "--last"):
                arg_last = True
            elif opt in ("-t", "--tor"):
                arg_use_tor = True

        if len(argv) == 0:
            usage()
            return
        else:
            target = argv[0]

        # If it's a file path, open a session on it.
        if arg_is_file:
            target = os.path.expanduser(target)

            if not os.path.exists(target) or not os.path.isfile(target):
                print_error("File not found")
                return

            __sessions__.new(target)
        # If it's a URL, download it and open a session on the temporary
        # file.
        elif arg_is_url:
            data = download(url=target, tor=arg_use_tor)

            if data:
                tmp = tempfile.NamedTemporaryFile(delete=False)
                tmp.write(data)
                tmp.close()

                __sessions__.new(tmp.name)
        # Try to open the specified file from the list of results from
        # the last find command.
        elif arg_last:
            if __sessions__.find:
                count = 1
                for item in __sessions__.find:
                    if count == int(target):
                        __sessions__.new(get_sample_path(item.sha256))
                        break

                    count += 1
            else:
                print_warning("You haven't performed a find yet")
        # Otherwise we assume it's an hash of an previously stored sample.
        else:
            target = argv[0].strip().lower()

            if len(target) == 32:
                key = "md5"
            elif len(target) == 64:
                key = "sha256"
            else:
                usage()
                return

            rows = self.db.find(key=key, value=target)

            if not rows:
                print_warning("No file found with the given hash {0}".format(target))
                return

            path = get_sample_path(rows[0].sha256)
            if path:
                __sessions__.new(path)

    ##
    # CLOSE
    #
    # This command resets the open session.
    # After that, all handles to the opened file should be closed and the
    # shell should be restored to the default prompt.
    def cmd_close(self, *args):
        __sessions__.close()

    ##
    # INFO
    #
    # This command returns information on the open session. It returns details
    # on the file (e.g. hashes) and other information that might available from
    # the database.
    def cmd_info(self, *args):
        if __sessions__.is_set():
            print(
                table(
                    ["Key", "Value"],
                    [
                        ("Name", __sessions__.current.file.name),
                        ("Tags", __sessions__.current.file.tags),
                        ("Path", __sessions__.current.file.path),
                        ("Size", __sessions__.current.file.size),
                        ("Type", __sessions__.current.file.type),
                        ("Mime", __sessions__.current.file.mime),
                        ("MD5", __sessions__.current.file.md5),
                        ("SHA1", __sessions__.current.file.sha1),
                        ("SHA256", __sessions__.current.file.sha256),
                        ("SHA512", __sessions__.current.file.sha512),
                        ("SSdeep", __sessions__.current.file.ssdeep),
                        ("CRC32", __sessions__.current.file.crc32),
                    ],
                )
            )

    ##
    # NOTES
    #
    # This command allows you to view, add, modify and delete notes associated
    # with the currently opened file.
    def cmd_notes(self, *args):
        def usage():
            print("usage: notes [-h] [-l] [-a] [-e <note id>] [-d <note id>]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--list (-h)\tList all notes available for the current file")
            print("\t--add (-a)\tAdd a new note to the current file")
            print("\t--view (-v)\tView the specified note")
            print("\t--edit (-e)\tEdit an existing note")
            print("\t--delete (-d)\tDelete an existing note")
            print("")

        try:
            opts, argv = getopt.getopt(args, "hlav:e:d:", ["help", "list", "add", "view=", "edit=", "delete="])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_list = False
        arg_add = False
        arg_view = None
        arg_edit = None
        arg_delete = None

        for opt, value in opts:
            if opt in ("-h", "--help"):
                help()
                return
            elif opt in ("-l", "--list"):
                arg_list = True
            elif opt in ("-a", "--add"):
                arg_add = True
            elif opt in ("-v", "--view"):
                arg_view = value
            elif opt in ("-e", "--edit"):
                arg_edit = value
            elif opt in ("-d", "--delete"):
                arg_delete = value

        if not __sessions__.is_set():
            print_error("No session opened")
            return

        if arg_list:
            # Retrieve all notes for the currently opened file.
            notes = Database().find(key="sha256", value=__sessions__.current.file.sha256)[0].note
            if not notes:
                print_info("No notes available for this file yet")
                return

            # Populate table rows.
            rows = []
            for note in notes:
                rows.append([note.id, note.title])

            # Display list of existing notes.
            print(table(header=["ID", "Title"], rows=rows))

        elif arg_add:
            title = raw_input("Enter a title for the new note: ")

            # Create a new temporary file.
            tmp = tempfile.NamedTemporaryFile(delete=False)
            # Open the temporary file with the default editor, or with nano.
            os.system('"${EDITOR:-nano}" ' + tmp.name)
            # Once the user is done editing, we need to read the content and
            # store it in the database.
            body = tmp.read()
            Database().add_note(__sessions__.current.file.sha256, title, body)
            # Finally, remove the temporary file.
            os.remove(tmp.name)

            print_info('New note with title "{0}" added to the current file'.format(bold(title)))

        elif arg_view:
            # Retrieve note wth the specified ID and print it.
            note = Database().get_note(arg_view)
            if note:
                print_info(bold("Title: ") + note.title)
                print_info(bold("Body:"))
                print(note.body)
            else:
                print_info("There is no note with ID {0}".format(arg_view))

        elif arg_edit:
            # Retrieve note with the specified ID.
            note = Database().get_note(arg_edit)
            if note:
                # Create a new temporary file.
                tmp = tempfile.NamedTemporaryFile(delete=False)
                # Write the old body to the temporary file.
                tmp.write(note.body)
                tmp.close()
                # Open the old body with the text editor.
                os.system('"${EDITOR:-nano}" ' + tmp.name)
                # Read the new body from the temporary file.
                body = open(tmp.name, "r").read()
                # Update the note entry with the new body.
                Database().edit_note(arg_edit, body)
                # Remove the temporary file.
                os.remove(tmp.name)

                print_info("Updated note with ID {0}".format(arg_edit))

        elif arg_delete:
            # Delete the note with the specified ID.
            Database().delete_note(arg_delete)
        else:
            usage()

    ##
    # STORE
    #
    # This command stores the opened file in the local repository and tries
    # to store details in the database.
    def cmd_store(self, *args):
        def usage():
            print("usage: store [-h] [-d] [-f <path>] [-s <size>] [-y <type>] [-n <name>] [-t]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--delete (-d)\tDelete the original file")
            print("\t--folder (-f)\tSpecify a folder to import")
            print("\t--file-size (-s)\tSpecify a maximum file size")
            print("\t--file-type (-y)\tSpecify a file type pattern")
            print("\t--file-name (-n)\tSpecify a file name pattern")
            print("\t--tags (-t)\tSpecify a list of comma-separated tags")
            print("")

        try:
            opts, argv = getopt.getopt(
                args, "hdf:s:y:n:t:", ["help", "delete", "folder=", "file-size=", "file-type=", "file-name=", "tags="]
            )
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_delete = False
        arg_folder = False
        arg_file_size = None
        arg_file_type = None
        arg_file_name = None
        arg_tags = None

        for opt, value in opts:
            if opt in ("-h", "--help"):
                help()
                return
            elif opt in ("-d", "--delete"):
                arg_delete = True
            elif opt in ("-f", "--folder"):
                arg_folder = value
            elif opt in ("-s", "--file-size"):
                arg_file_size = value
            elif opt in ("-y", "--file-type"):
                arg_file_type = value
            elif opt in ("-n", "--file-name"):
                arg_file_name = value
            elif opt in ("-t", "--tags"):
                arg_tags = value

        def add_file(obj, tags=None):
            if get_sample_path(obj.sha256):
                print_warning('Skip, file "{0}" appears to be already stored'.format(obj.name))
                return False

            # Store file to the local repository.
            new_path = store_sample(obj)
            if new_path:
                # Add file to the database.
                status = self.db.add(obj=obj, tags=tags)
                print_success('Stored file "{0}" to {1}'.format(obj.name, new_path))

            # Delete the file if requested to do so.
            if arg_delete:
                try:
                    os.unlink(obj.path)
                except Exception as e:
                    print_warning("Failed deleting file: {0}".format(e))

            return True

        # If the user specified the --folder flag, we walk recursively and try
        # to add all contained files to the local repository.
        # This is note going to open a new session.
        # TODO: perhaps disable or make recursion optional?
        if arg_folder:
            # Check if the specified folder is valid.
            if os.path.isdir(arg_folder):
                # Walk through the folder and subfolders.
                for dir_name, dir_names, file_names in os.walk(arg_folder):
                    # Add each collected file.
                    for file_name in file_names:
                        file_path = os.path.join(dir_name, file_name)

                        if not os.path.exists(file_path):
                            continue
                        # Check if file is not zero
                        if not os.path.getsize(file_path) > 0:
                            continue

                        # Check if the file name matches the provided pattern.
                        if arg_file_name:
                            if not fnmatch.fnmatch(file_name, arg_file_name):
                                # print_warning("Skip, file \"{0}\" doesn't match the file name pattern".format(file_path))
                                continue

                        # Check if the file type matches the provided pattern.
                        if arg_file_type:
                            if arg_file_type not in File(file_path).type:
                                # print_warning("Skip, file \"{0}\" doesn't match the file type".format(file_path))
                                continue

                        # Check if file exceeds maximum size limit.
                        if arg_file_size:
                            # Obtain file size.
                            if os.path.getsize(file_path) > arg_file_size:
                                print_warning('Skip, file "{0}" is too big'.format(file_path))
                                continue

                        file_obj = File(file_path)

                        # Add file.
                        add_file(file_obj, arg_tags)
            else:
                print_error("You specified an invalid folder: {0}".format(arg_folder))
        # Otherwise we try to store the currently opened file, if there is any.
        else:
            if __sessions__.is_set():
                # Add file.
                if add_file(__sessions__.current.file, arg_tags):
                    # Open session to the new file.
                    self.cmd_open(*[__sessions__.current.file.sha256])
            else:
                print_error("No session opened")

    ##
    # DELETE
    #
    # This commands deletes the currenlty opened file (only if it's stored in
    # the local repository) and removes the details from the database
    def cmd_delete(self, *args):
        if __sessions__.is_set():
            while True:
                choice = raw_input("Are you sure you want to delete this binary? Can't be reverted! [y/n] ")
                if choice == "y":
                    break
                elif choice == "n":
                    return

            rows = self.db.find("sha256", __sessions__.current.file.sha256)
            if rows:
                malware_id = rows[0].id
                if self.db.delete(malware_id):
                    print_success("File deleted")
                else:
                    print_error("Unable to delete file")

            os.remove(__sessions__.current.file.path)
            __sessions__.close()
        else:
            print_error("No session opened")

    ##
    # FIND
    #
    # This command is used to search for files in the database.
    def cmd_find(self, *args):
        def usage():
            print("usage: find [-h] [-t] <all|latest|name|md5|sha256|tag|note> <value>")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--tags (-t)\tList tags")
            print("")

        try:
            opts, argv = getopt.getopt(args, "ht", ["help", "tags"])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_list_tags = False

        for opt, value in opts:
            if opt in ("-h", "--help"):
                help()
                return
            elif opt in ("-t", "--tags"):
                arg_list_tags = True

        # One of the most useful search terms is by tag. With the --tags
        # argument we first retrieve a list of existing tags and the count
        # of files associated with each of them.
        if arg_list_tags:
            # Retrieve list of tags.
            tags = self.db.list_tags()

            if tags:
                rows = []
                # For each tag, retrieve the count of files associated with it.
                for tag in tags:
                    count = len(self.db.find("tag", tag.tag))
                    rows.append([tag.tag, count])

                # Generate the table with the results.
                header = ["Tag", "# Entries"]
                print(table(header=header, rows=rows))
            else:
                print("No tags available")

            return

        # At this point, if there are no search terms specified, return.
        if len(args) == 0:
            usage()
            return

        # The first argument is the search term (or "key").
        key = args[0]
        try:
            # The second argument is the search value.
            value = args[1]
        except IndexError:
            # If the user didn't specify any value, just set it to None.
            # Mostly pointless for now, but might have some useful cases
            # in the future.
            value = None

        # Search all the files matching the given parameters.
        items = self.db.find(key, value)
        if not items:
            return

        # Populate the list of search results.
        rows = []
        count = 1
        for item in items:
            rows.append([count, item.name, item.mime, item.md5])
            count += 1

        # Update find results in current session.
        __sessions__.find = items

        # Generate a table with the results.
        print(table(["#", "Name", "Mime", "MD5"], rows))

    ##
    # TAGS
    #
    # This command is used to modify the tags of the opened file.
    def cmd_tags(self, *args):
        def usage():
            print("usage: tags [-h] [-a=tags] [-d=tag]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--add (-a)\tAdd tags to the opened file (comma separated)")
            print("\t--delete (-d)\tDelete a tag from the opened file")
            print("")

        try:
            opts, argv = getopt.getopt(args, "ha:d:", ["help", "add=", "delete="])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_add = None
        arg_delete = None

        for opt, value in opts:
            if opt in ("-h", "--help"):
                help()
                return
            elif opt in ("-a", "--add"):
                arg_add = value
            elif opt in ("-d", "--delete"):
                arg_delete = value

        # This command requires a session to be opened.
        if not __sessions__.is_set():
            print_error("No session opened")
            return

        # If no arguments are specified, there's not much to do.
        # However, it could make sense to also retrieve a list of existing
        # tags from this command, and not just from the "find" command alone.
        if not arg_add and not arg_delete:
            usage()
            return

        # TODO: handle situation where addition or deletion of a tag fail.

        if arg_add:
            # Add specified tags to the database's entry belonging to
            # the opened file.
            db = Database()
            db.add_tags(__sessions__.current.file.sha256, arg_add)
            print_info("Tags added to the currently opened file")

            # We refresh the opened session to update the attributes.
            # Namely, the list of tags returned by the "info" command
            # needs to be re-generated, or it wouldn't show the new tags
            # until the existing session is closed a new one is opened.
            print_info("Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

        if arg_delete:
            # Delete the tag from the database.
            Database().delete_tag(arg_delete)
            # Refresh the session so that the attributes of the file are
            # updated.
            print_info("Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

    ###
    # SESSION
    #
    # This command is used to list and switch across all the opened sessions.
    def cmd_sessions(self, *args):
        def usage():
            print("usage: sessions [-h] [-l] [-s=session]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--list (-l)\tList all existing sessions")
            print("\t--switch (-s)\tSwitch to the specified session")
            print("")

        try:
            opts, argv = getopt.getopt(args, "hls:", ["help", "list", "switch="])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_list = False
        arg_switch = None

        for opt, value in opts:
            if opt in ("-h", "--help"):
                help()
                return
            elif opt in ("-l", "--list"):
                arg_list = True
            elif opt in ("-s", "--switch"):
                arg_switch = int(value)

        if arg_list:
            if not __sessions__.sessions:
                print_info("There are no opened sessions")
                return

            rows = []
            for session in __sessions__.sessions:
                current = ""
                if session == __sessions__.current:
                    current = "Yes"

                rows.append([session.id, session.file.name, session.file.md5, session.created_at, current])

            print_info("Opened Sessions:")
            print(table(header=["#", "Name", "MD5", "Created At", "Current"], rows=rows))
            return
        elif arg_switch:
            for session in __sessions__.sessions:
                if arg_switch == session.id:
                    __sessions__.switch(session)
                    return

            print_warning("The specified session ID doesn't seem to exist")
            return

        usage()

    ##
    # PROJECTS
    #
    # This command retrieves a list of all projects.
    # You can also switch to a different project.
    def cmd_projects(self, *args):
        def usage():
            print("usage: projects [-h] [-l] [-s=project]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--list (-l)\tList all existing projects")
            print("\t--switch (-s)\tSwitch to the specified project")
            print("")

        try:
            opts, argv = getopt.getopt(args, "hls:", ["help", "list", "switch="])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_list = False
        arg_switch = None

        for opt, value in opts:
            if opt in ("-h", "--help"):
                help()
                return
            elif opt in ("-l", "--list"):
                arg_list = True
            elif opt in ("-s", "--switch"):
                arg_switch = value

        projects_path = os.path.join(os.getcwd(), "projects")

        if not os.path.exists(projects_path):
            print_info("The projects directory does not exist yet")
            return

        if arg_list:
            print_info("Projects Available:")

            rows = []
            for project in os.listdir(projects_path):
                project_path = os.path.join(projects_path, project)
                if os.path.isdir(project_path):
                    current = ""
                    if __project__.name and project == __project__.name:
                        current = "Yes"
                    rows.append([project, time.ctime(os.path.getctime(project_path)), current])

            print(table(header=["Project Name", "Creation Time", "Current"], rows=rows))
            return
        elif arg_switch:
            if __sessions__.is_set():
                __sessions__.close()
                print_info("Closed opened session")

            __project__.open(arg_switch)
            print_info("Switched to project {0}".format(bold(arg_switch)))

            # Need to re-initialize the Database to open the new SQLite file.
            self.db = Database()
            return

        usage()
示例#18
0
文件: commands.py 项目: RATBORG/viper
class Commands(object):
    def __init__(self):
        # Open connection to the database.
        self.db = Database()

        # Map commands to their related functions.
        self.commands = dict(
            help=dict(obj=self.cmd_help, description="Show this help message"),
            open=dict(obj=self.cmd_open, description="Open a file"),
            close=dict(obj=self.cmd_close, description="Close the current session"),
            info=dict(obj=self.cmd_info, description="Show information on the opened file"),
            clear=dict(obj=self.cmd_clear, description="Clear the console"),
            store=dict(obj=self.cmd_store, description="Store the opened file to the local repository"),
            delete=dict(obj=self.cmd_delete, description="Delete the opened file"),
            find=dict(obj=self.cmd_find, description="Find a file"),
            tags=dict(obj=self.cmd_tags, description="Modify tags of the opened file"),
        )

    ##
    # CLEAR
    #
    # This command simply clears the shell.
    def cmd_clear(self, *args):
        os.system("clear")

    ##
    # HELP
    #
    # This command simply prints the help message.
    # It lists both embedded commands and loaded modules.
    def cmd_help(self, *args):
        print(bold("Commands:"))

        rows = []
        for command_name, command_item in self.commands.items():
            rows.append([command_name, command_item["description"]])

        print(table(["Command", "Description"], rows))
        print("")
        print(bold("Modules:"))

        rows = []
        for module_name, module_item in __modules__.items():
            rows.append([module_name, module_item["description"]])

        print(table(["Command", "Description"], rows))

    ##
    # OPEN
    #
    # This command is used to open a session on a given file.
    # It either can be an external file path, or a SHA256 hash of a file which
    # has been previously imported and stored.
    # While the session is active, every operation and module executed will be
    # run against the file specified.
    def cmd_open(self, *args):
        def usage():
            print("usage: open [-h] [-f] [-u] [-t] <target>")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--file (-f)\tThe target is a file")
            print("\t--url (-u)\tThe target is a URL")
            print("\t--tor (-t)\tDownload the file through Tor")
            print("")
            print("You can also specify a SHA256 hash to a previously stored")
            print("file in order to open a session on it.")
            print("")

        try:
            opts, argv = getopt.getopt(args, "hfut", ["help", "file", "url", "tor"])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_is_file = False
        arg_is_url = False
        arg_use_tor = False

        for opt, value in opts:
            if opt in ("-h", "--help"):
                help()
                return
            elif opt in ("-f", "--file"):
                arg_is_file = True
            elif opt in ("-u", "--url"):
                arg_is_url = True
            elif opt in ("-t", "--tor"):
                arg_use_tor = True

        if len(argv) == 0:
            usage()
            return
        else:
            target = argv[0]

        # If it's a file path, open a session on it.
        if arg_is_file:
            target = os.path.expanduser(target)

            if not os.path.exists(target) or not os.path.isfile(target):
                print_error("File not found")
                return

            __session__.set(target)
        # If it's a URL, download it and open a session on the temporary
        # file.
        elif arg_is_url:
            data = download(url=target, tor=arg_use_tor)

            if data:
                tmp = tempfile.NamedTemporaryFile(delete=False)
                tmp.write(data)
                tmp.close()

                __session__.set(tmp.name)
        # Otherwise we assume it's an hash of an previously stored sample.
        else:
            target = argv[0].strip().lower()
            path = get_sample_path(target)
            if path:
                __session__.set(path)

    ##
    # CLOSE
    #
    # This command resets the open session.
    # After that, all handles to the opened file should be closed and the
    # shell should be restored to the default prompt.
    def cmd_close(self, *args):
        __session__.clear()

    ##
    # INFO
    #
    # This command returns information on the open session. It returns details
    # on the file (e.g. hashes) and other information that might available from
    # the database.
    def cmd_info(self, *args):
        if __session__.is_set():
            print(
                table(
                    ["Key", "Value"],
                    [
                        ("Name", __session__.file.name),
                        ("Tags", __session__.file.tags),
                        ("Path", __session__.file.path),
                        ("Size", __session__.file.size),
                        ("Type", __session__.file.type),
                        ("MD5", __session__.file.md5),
                        ("SHA1", __session__.file.sha1),
                        ("SHA256", __session__.file.sha256),
                        ("SHA512", __session__.file.sha512),
                        ("SSdeep", __session__.file.ssdeep),
                        ("CRC32", __session__.file.crc32),
                    ],
                )
            )

    ##
    # STORE
    #
    # This command stores the opened file in the local repository and tries
    # to store details in the database.
    def cmd_store(self, *args):
        def usage():
            print("usage: store [-h] [-d] [-f <path>] [-s <size>] [-y <type>] [-n <name>] [-t]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--delete (-d)\tDelete the original file")
            print("\t--folder (-f)\tSpecify a folder to import")
            print("\t--file-size (-s)\tSpecify a maximum file size")
            print("\t--file-type (-y)\tSpecify a file type pattern")
            print("\t--file-name (-n)\tSpecify a file name pattern")
            print("\t--tags (-t)\tSpecify a list of comma-separated tags")
            print("")

        try:
            opts, argv = getopt.getopt(
                args, "hdf:s:y:n:t:", ["help", "delete", "folder=", "file-size=", "file-type=", "file-name=", "tags="]
            )
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_delete = False
        arg_folder = False
        arg_file_size = None
        arg_file_type = None
        arg_file_name = None
        arg_tags = None

        for opt, value in opts:
            if opt in ("-h", "--help"):
                help()
                return
            elif opt in ("-d", "--delete"):
                arg_delete = True
            elif opt in ("-f", "--folder"):
                arg_folder = value
            elif opt in ("-s", "--file-size"):
                arg_file_size = value
            elif opt in ("-y", "--file-type"):
                arg_file_type = value
            elif opt in ("-n", "--file-name"):
                arg_file_name = value
            elif opt in ("-t", "--tags"):
                arg_tags = value

        def add_file(obj, tags=None):
            if get_sample_path(obj.sha256):
                print_warning('Skip, file "{0}" appears to be already stored'.format(obj.name))
                return False

            # Store file to the local repository.
            new_path = store_sample(obj)
            if new_path:
                # Add file to the database.
                status = self.db.add(obj=obj, tags=tags)
                print_success('Stored file "{0}" to {1}'.format(obj.name, new_path))

            # Delete the file if requested to do so.
            if arg_delete:
                try:
                    os.unlink(obj.path)
                except Exception as e:
                    print_warning("Failed deleting file: {0}".format(e))

            return True

        # If the user specified the --folder flag, we walk recursively and try
        # to add all contained files to the local repository.
        # This is note going to open a new session.
        # TODO: perhaps disable or make recursion optional?
        if arg_folder:
            # Check if the specified folder is valid.
            if os.path.isdir(arg_folder):
                # Walk through the folder and subfolders.
                for dir_name, dir_names, file_names in os.walk(arg_folder):
                    # Add each collected file.
                    for file_name in file_names:
                        file_path = os.path.join(dir_name, file_name)

                        if not os.path.exists(file_path):
                            continue

                        # Check if the file name matches the provided pattern.
                        if arg_file_name:
                            if not fnmatch.fnmatch(file_name, arg_file_name):
                                # print_warning("Skip, file \"{0}\" doesn't match the file name pattern".format(file_path))
                                continue

                        # Check if the file type matches the provided pattern.
                        if arg_file_type:
                            if arg_file_type not in File(file_path).type:
                                # print_warning("Skip, file \"{0}\" doesn't match the file type".format(file_path))
                                continue

                        # Check if file exceeds maximum size limit.
                        if arg_file_size:
                            # Obtain file size.
                            if os.path.getsize(file_path) > arg_file_size:
                                print_warning('Skip, file "{0}" is too big'.format(file_path))
                                continue

                        file_obj = File(file_path)

                        # Add file.
                        add_file(file_obj, arg_tags)
            else:
                print_error("You specified an invalid folder: {0}".format(arg_folder))
        # Otherwise we try to store the currently opened file, if there is any.
        else:
            if __session__.is_set():
                # Add file.
                if add_file(__session__.file, arg_tags):
                    # Open session to the new file.
                    self.cmd_open(*[__session__.file.sha256])
            else:
                print_error("No session opened")

    ##
    # DELETE
    #
    # This commands deletes the currenlty opened file (only if it's stored in
    # the local repository) and removes the details from the database
    def cmd_delete(self, *args):
        if __session__.is_set():
            while True:
                choice = raw_input("Are you sure you want to delete this binary? Can't be reverted! [y/n] ")
                if choice == "y":
                    break
                elif choice == "n":
                    return

            rows = self.db.find("sha256", __session__.file.sha256)
            if rows:
                malware_id = rows[0].id
                if self.db.delete(malware_id):
                    print_success("File deleted")
                else:
                    print_error("Unable to delete file")

            os.remove(__session__.file.path)
            __session__.clear()
        else:
            print_error("No session opened")

    ##
    # FIND
    #
    # This command is used to search for files in the database.
    def cmd_find(self, *args):
        def usage():
            print("usage: find [-h] [-t] <name|md5|sha256|tag> <value>")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--tags (-t)\tList tags")
            print("")

        try:
            opts, argv = getopt.getopt(args, "ht", ["help", "tags"])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_list_tags = False

        for opt, value in opts:
            if opt in ("-h", "--help"):
                help()
                return
            elif opt in ("-t", "--tags"):
                arg_list_tags = True

        # One of the most useful search terms is by tag. With the --tags
        # argument we first retrieve a list of existing tags and the count
        # of files associated with each of them.
        if arg_list_tags:
            # Retrieve list of tags.
            tags = self.db.list_tags()

            if tags:
                rows = []
                # For each tag, retrieve the count of files associated with it.
                for tag in tags:
                    count = len(self.db.find("tag", tag.tag))
                    rows.append([tag.tag, count])

                # Generate the table with the results.
                header = ["Tag", "# Entries"]
                print(table(header=header, rows=rows))
            else:
                print("No tags available")

            return

        # At this point, if there are no search terms specified, return.
        if len(args) == 0:
            usage()
            return

        # The first argument is the search term (or "key").
        key = args[0]
        try:
            # The second argument is the search value.
            value = args[1]
        except IndexError:
            # If the user didn't specify any value, just set it to None.
            # Mostly pointless for now, but might have some useful cases
            # in the future.
            value = None

        # Search all the files matching the given parameters.
        items = self.db.find(key, value)
        if not items:
            return

        # Populate the list of search results.
        rows = []
        for item in items:
            rows.append([item.name, item.type, item.sha256])

        # Generate a table with the results.
        print(table(["Name", "Type", "SHA256"], rows))

    ##
    # TAGS
    #
    # This command is used to modify the tags of the opened file.
    def cmd_tags(self, *args):
        def usage():
            print("usage: tags [-h] [-a=tags] [-d=tag]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--add (-a)\tAdd tags to the opened file (comma separated)")
            print("\t--delete (-d)\tDelete a tag from the opened file")
            print("")

        try:
            opts, argv = getopt.getopt(args, "ha:d:", ["help", "add=", "delete="])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_add = None
        arg_delete = None

        for opt, value in opts:
            if opt in ("-h", "--help"):
                help()
                return
            elif opt in ("-a", "--add"):
                arg_add = value
            elif opt in ("-d", "--delete"):
                arg_delete = value

        # This command requires a session to be opened.
        if not __session__.is_set():
            print_error("No session opened")
            return

        # If no arguments are specified, there's not much to do.
        # However, it could make sense to also retrieve a list of existing
        # tags from this command, and not just from the "find" command alone.
        if not arg_add and not arg_delete:
            print_error("You need to specify an option, either add or delete")
            return

        if arg_add:
            # Add specified tags to the database's entry belonging to
            # the opened file.
            db = Database()
            db.add_tags(__session__.file.sha256, arg_add)
            print_info("Tags added to the currently opened file")

            # We refresh the opened session to update the attributes.
            # Namely, the list of tags returned by the "info" command
            # needs to be re-generated, or it wouldn't show the new tags
            # until the existing session is closed a new one is opened.
            print_info("Refreshing session to update attributes...")
            __session__.set(__session__.file.path)

        if arg_delete:
            # TODO
            pass
示例#19
0
文件: web.py 项目: lksnyder0/viper
def add_file():
    tags = request.forms.get('tag_list')
    uploads = request.files.getlist('file')

    # Set Project
    project = request.forms.get('project')
    if project in project_list():
        __project__.open(project)
    else:
        __project__.open('../')
        project = 'Main'
    db = Database()
    file_list = []
    # Write temp file to disk
    with upload_temp() as temp_dir:
        for upload in uploads:
            file_path = os.path.join(temp_dir, upload.filename)
            with open(file_path, 'w') as tmp_file:
                tmp_file.write(upload.file.read())
            # Zip Files
            if request.forms.get('compression') == 'zip':
                zip_pass = request.forms.get('zip_pass')
                try:
                    with ZipFile(file_path) as zf:
                        zf.extractall(temp_dir, pwd=zip_pass)
                    for root, dirs, files in os.walk(temp_dir, topdown=False):
                        for name in files:
                            if not name == upload.filename:
                                file_list.append(os.path.join(root, name))
                except Exception as e:
                    return template('error.tpl',
                                    error="Error with zipfile - {0}".format(e))
            # GZip Files
            elif request.forms.get('compression') == 'gz':
                try:
                    gzf = GzipFile(file_path, 'rb')
                    decompress = gzf.read()
                    gzf.close()
                    with open(file_path[:-3], "wb") as df:
                        df.write(decompress)
                    file_list.append(file_path[:-3])
                except Exception as e:
                    return template(
                        'error.tpl',
                        error="Error with gzipfile - {0}".format(e))
            # BZip2 Files
            elif request.forms.get('compression') == 'bz2':
                try:
                    bz2f = BZ2File(file_path, 'rb')
                    decompress = bz2f.read()
                    bz2f.close()
                    with open(file_path[:-3], "wb") as df:
                        df.write(decompress)
                    file_list.append(file_path[:-3])
                except Exception as e:
                    return template(
                        'error.tpl',
                        error="Error with bzip2file - {0}".format(e))
            # Tar Files (any, including tar.gz tar.bz2)
            elif request.forms.get('compression') == 'tar':
                try:
                    if not tarfile.is_tarfile(file_path):
                        return template('error.tpl',
                                        error="This is not a tar file")
                    with tarfile.open(file_path, 'r:*') as tarf:
                        tarf.extractall(temp_dir)
                    for root, dirs, files in os.walk(temp_dir, topdown=False):
                        for name in files:
                            if not name == upload.filename:
                                file_list.append(os.path.join(root, name))
                except Exception as e:
                    return template('error.tpl',
                                    error="Error with tarfile - {0}".format(e))
            # Non zip files
            elif request.forms.get('compression') == 'none':
                file_list.append(file_path)

        # Add each file
        for new_file in file_list:
            print new_file
            obj = File(new_file)
            new_path = store_sample(obj)
            success = True
            if new_path:
                # Add file to the database.
                success = db.add(obj=obj, tags=tags)
                if not success:
                    return template(
                        'error.tpl',
                        error="Unable to Store The File: {0}".format(
                            upload.filename))
    redirect("/project/{0}".format(project))
示例#20
0
def add_file():
    tags = request.forms.get('tag_list')
    uploads = request.files.getlist('file')
    
    # Set Project
    project = request.forms.get('project')
    if project in project_list():
        __project__.open(project)
    else:
        __project__.open('../')
        project = 'Main'
    db = Database()    
    file_list = []
    # Write temp file to disk
    with upload_temp() as temp_dir:
        for upload in uploads:
            file_path = os.path.join(temp_dir, upload.filename)
            with open(file_path, 'w') as tmp_file:
                tmp_file.write(upload.file.read())
            # Zip Files
            if request.forms.get('compression') == 'zip':
                zip_pass = request.forms.get('zip_pass')
                try:
                    with ZipFile(file_path) as zf:
                        zf.extractall(temp_dir, pwd=zip_pass)            
                    for root, dirs, files in os.walk(temp_dir, topdown=False):
                        for name in files:
                            if not name == upload.filename:
                                file_list.append(os.path.join(root, name))
                except Exception as e:
                    return template('error.tpl', error="Error with zipfile - {0}".format(e))
            # GZip Files
            elif request.forms.get('compression') == 'gz':
                try:
                    gzf = GzipFile(file_path, 'rb')
                    decompress = gzf.read()
                    gzf.close()
                    with open(file_path[:-3],"wb") as df:
                        df.write(decompress)
                    file_list.append(file_path[:-3])
                except Exception as e:
                    return template('error.tpl', error="Error with gzipfile - {0}".format(e))
            # BZip2 Files
            elif request.forms.get('compression') == 'bz2':
                try:
                    bz2f = BZ2File(file_path, 'rb')
                    decompress = bz2f.read()
                    bz2f.close()
                    with open(file_path[:-3],"wb") as df:
                        df.write(decompress)
                    file_list.append(file_path[:-3])
                except Exception as e:
                    return template('error.tpl', error="Error with bzip2file - {0}".format(e))
            # Tar Files (any, including tar.gz tar.bz2)
            elif request.forms.get('compression') == 'tar':
                try:
                    if not tarfile.is_tarfile(file_path):
                        return template('error.tpl', error="This is not a tar file")
                    with tarfile.open(file_path,'r:*') as tarf:
                        tarf.extractall(temp_dir)
                    for root, dirs, files in os.walk(temp_dir, topdown=False):
                        for name in files:
                            if not name == upload.filename:
                                file_list.append(os.path.join(root, name))
                except Exception as e:
                    return template('error.tpl', error="Error with tarfile - {0}".format(e))
            # Non zip files
            elif request.forms.get('compression') == 'none':
                file_list.append(file_path)
            
        # Add each file
        for new_file in file_list:
            print new_file
            obj = File(new_file)
            new_path = store_sample(obj)
            success = True
            if new_path:
                # Add file to the database.
                success = db.add(obj=obj, tags=tags)
                if not success:
                    return template('error.tpl', error="Unable to Store The File: {0}".format(upload.filename))
    redirect("/project/{0}".format(project))
示例#21
0
文件: commands.py 项目: vicgc/viper
class Commands(object):
    def __init__(self):
        # Open connection to the database.
        self.db = Database()

        # Map commands to their related functions.
        self.commands = dict(
            help=dict(obj=self.cmd_help, description="Show this help message"),
            open=dict(obj=self.cmd_open, description="Open a file"),
            close=dict(obj=self.cmd_close,
                       description="Close the current session"),
            info=dict(obj=self.cmd_info,
                      description="Show information on the opened file"),
            notes=dict(
                obj=self.cmd_notes,
                description="View, add and edit notes on the opened file"),
            clear=dict(obj=self.cmd_clear, description="Clear the console"),
            store=dict(
                obj=self.cmd_store,
                description="Store the opened file to the local repository"),
            delete=dict(obj=self.cmd_delete,
                        description="Delete the opened file"),
            find=dict(obj=self.cmd_find, description="Find a file"),
            tags=dict(obj=self.cmd_tags,
                      description="Modify tags of the opened file"),
            sessions=dict(obj=self.cmd_sessions,
                          description="List or switch sessions"),
            projects=dict(obj=self.cmd_projects,
                          description="List or switch existing projects"),
            export=dict(
                obj=self.cmd_export,
                description="Export the current session to file or zip"),
        )

    ##
    # CLEAR
    #
    # This command simply clears the shell.
    def cmd_clear(self, *args):
        os.system('clear')

    ##
    # HELP
    #
    # This command simply prints the help message.
    # It lists both embedded commands and loaded modules.
    def cmd_help(self, *args):
        print(bold("Commands:"))

        rows = []
        for command_name, command_item in self.commands.items():
            rows.append([command_name, command_item['description']])

        rows = sorted(rows, key=lambda entry: entry[0])

        print(table(['Command', 'Description'], rows))
        print("")
        print(bold("Modules:"))

        rows = []
        for module_name, module_item in __modules__.items():
            rows.append([module_name, module_item['description']])

        rows = sorted(rows, key=lambda entry: entry[0])

        print(table(['Command', 'Description'], rows))

    ##
    # OPEN
    #
    # This command is used to open a session on a given file.
    # It either can be an external file path, or a SHA256 hash of a file which
    # has been previously imported and stored.
    # While the session is active, every operation and module executed will be
    # run against the file specified.
    def cmd_open(self, *args):
        def usage():
            print("usage: open [-h] [-f] [-u] [-l] [-t] <target|md5|sha256>")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--file (-f)\tThe target is a file")
            print("\t--url (-u)\tThe target is a URL")
            print(
                "\t--last (-l)\tThe target is the entry number from the last find command's results"
            )
            print("\t--tor (-t)\tDownload the file through Tor")
            print("")
            print(
                "You can also specify a MD5 or SHA256 hash to a previously stored"
            )
            print("file in order to open a session on it.")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hfult',
                                       ['help', 'file', 'url', 'last', 'tor'])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_is_file = False
        arg_is_url = False
        arg_last = False
        arg_use_tor = False

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-f', '--file'):
                arg_is_file = True
            elif opt in ('-u', '--url'):
                arg_is_url = True
            elif opt in ('-l', '--last'):
                arg_last = True
            elif opt in ('-t', '--tor'):
                arg_use_tor = True

        if len(argv) == 0:
            usage()
            return
        else:
            target = argv[0]

        # If it's a file path, open a session on it.
        if arg_is_file:
            target = os.path.expanduser(target)

            # This is kind of hacky. It checks if there are additional arguments
            # to the open command, if there is I assume that it's the continuation
            # of a filename with spaces. I then concatenate them.
            # TODO: improve this.
            if len(argv) > 1:
                for arg in argv[1:]:
                    target += ' ' + arg

            if not os.path.exists(target) or not os.path.isfile(target):
                print_error("File not found: {0}".format(target))
                return

            __sessions__.new(target)
        # If it's a URL, download it and open a session on the temporary
        # file.
        elif arg_is_url:
            data = download(url=target, tor=arg_use_tor)

            if data:
                tmp = tempfile.NamedTemporaryFile(delete=False)
                tmp.write(data)
                tmp.close()

                __sessions__.new(tmp.name)
        # Try to open the specified file from the list of results from
        # the last find command.
        elif arg_last:
            if __sessions__.find:
                count = 1
                for item in __sessions__.find:
                    if count == int(target):
                        __sessions__.new(get_sample_path(item.sha256))
                        break

                    count += 1
            else:
                print_warning("You haven't performed a find yet")
        # Otherwise we assume it's an hash of an previously stored sample.
        else:
            target = argv[0].strip().lower()

            if len(target) == 32:
                key = 'md5'
            elif len(target) == 64:
                key = 'sha256'
            else:
                usage()
                return

            rows = self.db.find(key=key, value=target)

            if not rows:
                print_warning(
                    "No file found with the given hash {0}".format(target))
                return

            path = get_sample_path(rows[0].sha256)
            if path:
                __sessions__.new(path)

    ##
    # CLOSE
    #
    # This command resets the open session.
    # After that, all handles to the opened file should be closed and the
    # shell should be restored to the default prompt.
    def cmd_close(self, *args):
        __sessions__.close()

    ##
    # INFO
    #
    # This command returns information on the open session. It returns details
    # on the file (e.g. hashes) and other information that might available from
    # the database.
    def cmd_info(self, *args):
        if __sessions__.is_set():
            print(
                table(['Key', 'Value'],
                      [('Name', __sessions__.current.file.name),
                       ('Tags', __sessions__.current.file.tags),
                       ('Path', __sessions__.current.file.path),
                       ('Size', __sessions__.current.file.size),
                       ('Type', __sessions__.current.file.type),
                       ('Mime', __sessions__.current.file.mime),
                       ('MD5', __sessions__.current.file.md5),
                       ('SHA1', __sessions__.current.file.sha1),
                       ('SHA256', __sessions__.current.file.sha256),
                       ('SHA512', __sessions__.current.file.sha512),
                       ('SSdeep', __sessions__.current.file.ssdeep),
                       ('CRC32', __sessions__.current.file.crc32)]))

    ##
    # NOTES
    #
    # This command allows you to view, add, modify and delete notes associated
    # with the currently opened file.
    def cmd_notes(self, *args):
        def usage():
            print("usage: notes [-h] [-l] [-a] [-e <note id>] [-d <note id>]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print(
                "\t--list (-l)\tList all notes available for the current file")
            print("\t--add (-a)\tAdd a new note to the current file")
            print("\t--view (-v)\tView the specified note")
            print("\t--edit (-e)\tEdit an existing note")
            print("\t--delete (-d)\tDelete an existing note")
            print("")

        try:
            opts, argv = getopt.getopt(
                args, 'hlav:e:d:',
                ['help', 'list', 'add', 'view=', 'edit=', 'delete='])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_list = False
        arg_add = False
        arg_view = None
        arg_edit = None
        arg_delete = None

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-l', '--list'):
                arg_list = True
            elif opt in ('-a', '--add'):
                arg_add = True
            elif opt in ('-v', '--view'):
                arg_view = value
            elif opt in ('-e', '--edit'):
                arg_edit = value
            elif opt in ('-d', '--delete'):
                arg_delete = value

        if not __sessions__.is_set():
            print_error("No session opened")
            return

        if arg_list:
            # Retrieve all notes for the currently opened file.
            malware = Database().find(key='sha256',
                                      value=__sessions__.current.file.sha256)
            if not malware:
                print_error(
                    "The opened file doesn't appear to be in the database, have you stored it yet?"
                )
                return

            notes = malware[0].note
            if not notes:
                print_info("No notes available for this file yet")
                return

            # Populate table rows.
            rows = []
            for note in notes:
                rows.append([note.id, note.title])

            # Display list of existing notes.
            print(table(header=['ID', 'Title'], rows=rows))

        elif arg_add:
            title = raw_input("Enter a title for the new note: ")

            # Create a new temporary file.
            tmp = tempfile.NamedTemporaryFile(delete=False)
            # Open the temporary file with the default editor, or with nano.
            os.system('"${EDITOR:-nano}" ' + tmp.name)
            # Once the user is done editing, we need to read the content and
            # store it in the database.
            body = tmp.read()
            Database().add_note(__sessions__.current.file.sha256, title, body)
            # Finally, remove the temporary file.
            os.remove(tmp.name)

            print_info(
                "New note with title \"{0}\" added to the current file".format(
                    bold(title)))

        elif arg_view:
            # Retrieve note wth the specified ID and print it.
            note = Database().get_note(arg_view)
            if note:
                print_info(bold('Title: ') + note.title)
                print_info(bold('Body:'))
                print(note.body)
            else:
                print_info("There is no note with ID {0}".format(arg_view))

        elif arg_edit:
            # Retrieve note with the specified ID.
            note = Database().get_note(arg_edit)
            if note:
                # Create a new temporary file.
                tmp = tempfile.NamedTemporaryFile(delete=False)
                # Write the old body to the temporary file.
                tmp.write(note.body)
                tmp.close()
                # Open the old body with the text editor.
                os.system('"${EDITOR:-nano}" ' + tmp.name)
                # Read the new body from the temporary file.
                body = open(tmp.name, 'r').read()
                # Update the note entry with the new body.
                Database().edit_note(arg_edit, body)
                # Remove the temporary file.
                os.remove(tmp.name)

                print_info("Updated note with ID {0}".format(arg_edit))

        elif arg_delete:
            # Delete the note with the specified ID.
            Database().delete_note(arg_delete)
        else:
            usage()

    ##
    # STORE
    #
    # This command stores the opened file in the local repository and tries
    # to store details in the database.
    def cmd_store(self, *args):
        def usage():
            print(
                "usage: store [-h] [-d] [-f <path>] [-s <size>] [-y <type>] [-n <name>] [-t]"
            )

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--delete (-d)\tDelete the original file")
            print("\t--folder (-f)\tSpecify a folder to import")
            print("\t--file-size (-s)\tSpecify a maximum file size")
            print("\t--file-type (-y)\tSpecify a file type pattern")
            print("\t--file-name (-n)\tSpecify a file name pattern")
            print("\t--tags (-t)\tSpecify a list of comma-separated tags")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hdf:s:y:n:t:', [
                'help', 'delete', 'folder=', 'file-size=', 'file-type=',
                'file-name=', 'tags='
            ])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_delete = False
        arg_folder = False
        arg_file_size = None
        arg_file_type = None
        arg_file_name = None
        arg_tags = None

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-d', '--delete'):
                arg_delete = True
            elif opt in ('-f', '--folder'):
                arg_folder = value
            elif opt in ('-s', '--file-size'):
                arg_file_size = value
            elif opt in ('-y', '--file-type'):
                arg_file_type = value
            elif opt in ('-n', '--file-name'):
                arg_file_name = value
            elif opt in ('-t', '--tags'):
                arg_tags = value

        def add_file(obj, tags=None):
            if get_sample_path(obj.sha256):
                print_warning(
                    "Skip, file \"{0}\" appears to be already stored".format(
                        obj.name))
                return False

            # Try to store file object into database.
            status = self.db.add(obj=obj, tags=tags)
            if status:
                # If succeeds, store also in the local repository.
                # If something fails in the database (for example unicode strings)
                # we don't want to have the binary lying in the repository with no
                # associated database record.
                new_path = store_sample(obj)
                print_success("Stored file \"{0}\" to {1}".format(
                    obj.name, new_path))
            else:
                return False

            # Delete the file if requested to do so.
            if arg_delete:
                try:
                    os.unlink(obj.path)
                except Exception as e:
                    print_warning("Failed deleting file: {0}".format(e))

            return True

        # If the user specified the --folder flag, we walk recursively and try
        # to add all contained files to the local repository.
        # This is note going to open a new session.
        # TODO: perhaps disable or make recursion optional?
        if arg_folder:
            # Check if the specified folder is valid.
            if os.path.isdir(arg_folder):
                # Walk through the folder and subfolders.
                for dir_name, dir_names, file_names in os.walk(arg_folder):
                    # Add each collected file.
                    for file_name in file_names:
                        file_path = os.path.join(dir_name, file_name)

                        if not os.path.exists(file_path):
                            continue
                        # Check if file is not zero.
                        if not os.path.getsize(file_path) > 0:
                            continue

                        # Check if the file name matches the provided pattern.
                        if arg_file_name:
                            if not fnmatch.fnmatch(file_name, arg_file_name):
                                #print_warning("Skip, file \"{0}\" doesn't match the file name pattern".format(file_path))
                                continue

                        # Check if the file type matches the provided pattern.
                        if arg_file_type:
                            if arg_file_type not in File(file_path).type:
                                #print_warning("Skip, file \"{0}\" doesn't match the file type".format(file_path))
                                continue

                        # Check if file exceeds maximum size limit.
                        if arg_file_size:
                            # Obtain file size.
                            if os.path.getsize(file_path) > arg_file_size:
                                print_warning(
                                    "Skip, file \"{0}\" is too big".format(
                                        file_path))
                                continue

                        file_obj = File(file_path)

                        # Add file.
                        add_file(file_obj, arg_tags)
            else:
                print_error(
                    "You specified an invalid folder: {0}".format(arg_folder))
        # Otherwise we try to store the currently opened file, if there is any.
        else:
            if __sessions__.is_set():
                if __sessions__.current.file.size == 0:
                    print_warning(
                        "Skip, file \"{0}\" appears to be empty".format(
                            __sessions__.current.file.name))
                    return False

                # Add file.
                if add_file(__sessions__.current.file, arg_tags):
                    # Open session to the new file.
                    self.cmd_open(*[__sessions__.current.file.sha256])
            else:
                print_error("No session opened")

    ##
    # DELETE
    #
    # This commands deletes the currenlty opened file (only if it's stored in
    # the local repository) and removes the details from the database
    def cmd_delete(self, *args):
        if __sessions__.is_set():
            while True:
                choice = raw_input(
                    "Are you sure you want to delete this binary? Can't be reverted! [y/n] "
                )
                if choice == 'y':
                    break
                elif choice == 'n':
                    return

            rows = self.db.find('sha256', __sessions__.current.file.sha256)
            if rows:
                malware_id = rows[0].id
                if self.db.delete(malware_id):
                    print_success("File deleted")
                else:
                    print_error("Unable to delete file")

            os.remove(__sessions__.current.file.path)
            __sessions__.close()
        else:
            print_error("No session opened")

    ##
    # FIND
    #
    # This command is used to search for files in the database.
    def cmd_find(self, *args):
        def usage():
            print(
                "usage: find [-h] [-t] <all|latest|name|type|mime|md5|sha256|tag|note> <value>"
            )

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--tags (-t)\tList tags")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'ht', ['help', 'tags'])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_list_tags = False

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-t', '--tags'):
                arg_list_tags = True

        # One of the most useful search terms is by tag. With the --tags
        # argument we first retrieve a list of existing tags and the count
        # of files associated with each of them.
        if arg_list_tags:
            # Retrieve list of tags.
            tags = self.db.list_tags()

            if tags:
                rows = []
                # For each tag, retrieve the count of files associated with it.
                for tag in tags:
                    count = len(self.db.find('tag', tag.tag))
                    rows.append([tag.tag, count])

                # Generate the table with the results.
                header = ['Tag', '# Entries']
                rows.sort(key=lambda x: x[1], reverse=True)
                print(table(header=header, rows=rows))
            else:
                print("No tags available")

            return

        # At this point, if there are no search terms specified, return.
        if len(args) == 0:
            usage()
            return

        # The first argument is the search term (or "key").
        key = args[0]
        if key != 'all' and key != 'latest':
            try:
                # The second argument is the search value.
                value = args[1]
            except IndexError:
                print_error("You need to include a search term.")
                return
        else:
            value = None

        # Search all the files matching the given parameters.
        items = self.db.find(key, value)
        if not items:
            return

        # Populate the list of search results.
        rows = []
        count = 1
        for item in items:
            tag = ', '.join([t.tag for t in item.tag if t.tag])
            row = [count, item.name, item.mime, item.md5, tag]
            if key == 'latest':
                row.append(item.created_at)

            rows.append(row)
            count += 1

        # Update find results in current session.
        __sessions__.find = items

        # Generate a table with the results.
        header = ['#', 'Name', 'Mime', 'MD5', 'Tags']
        if key == 'latest':
            header.append('Created At')

        print(table(header=header, rows=rows))

    ##
    # TAGS
    #
    # This command is used to modify the tags of the opened file.
    def cmd_tags(self, *args):
        def usage():
            print("usage: tags [-h] [-a=tags] [-d=tag]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print(
                "\t--add (-a)\tAdd tags to the opened file (comma separated)")
            print("\t--delete (-d)\tDelete a tag from the opened file")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'ha:d:',
                                       ['help', 'add=', 'delete='])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_add = None
        arg_delete = None

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-a', '--add'):
                arg_add = value
            elif opt in ('-d', '--delete'):
                arg_delete = value

        # This command requires a session to be opened.
        if not __sessions__.is_set():
            print_error("No session opened")
            return

        # If no arguments are specified, there's not much to do.
        # However, it could make sense to also retrieve a list of existing
        # tags from this command, and not just from the "find" command alone.
        if not arg_add and not arg_delete:
            usage()
            return

        # TODO: handle situation where addition or deletion of a tag fail.

        if arg_add:
            # Add specified tags to the database's entry belonging to
            # the opened file.
            db = Database()
            db.add_tags(__sessions__.current.file.sha256, arg_add)
            print_info("Tags added to the currently opened file")

            # We refresh the opened session to update the attributes.
            # Namely, the list of tags returned by the "info" command
            # needs to be re-generated, or it wouldn't show the new tags
            # until the existing session is closed a new one is opened.
            print_info("Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

        if arg_delete:
            # Delete the tag from the database.
            Database().delete_tag(arg_delete)
            # Refresh the session so that the attributes of the file are
            # updated.
            print_info("Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

    ###
    # SESSION
    #
    # This command is used to list and switch across all the opened sessions.
    def cmd_sessions(self, *args):
        def usage():
            print("usage: sessions [-h] [-l] [-s=session]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--list (-l)\tList all existing sessions")
            print("\t--switch (-s)\tSwitch to the specified session")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hls:',
                                       ['help', 'list', 'switch='])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_list = False
        arg_switch = None

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-l', '--list'):
                arg_list = True
            elif opt in ('-s', '--switch'):
                arg_switch = int(value)

        if arg_list:
            if not __sessions__.sessions:
                print_info("There are no opened sessions")
                return

            rows = []
            for session in __sessions__.sessions:
                current = ''
                if session == __sessions__.current:
                    current = 'Yes'

                rows.append([
                    session.id, session.file.name, session.file.md5,
                    session.created_at, current
                ])

            print_info("Opened Sessions:")
            print(
                table(header=['#', 'Name', 'MD5', 'Created At', 'Current'],
                      rows=rows))
            return
        elif arg_switch:
            for session in __sessions__.sessions:
                if arg_switch == session.id:
                    __sessions__.switch(session)
                    return

            print_warning("The specified session ID doesn't seem to exist")
            return

        usage()

    ##
    # PROJECTS
    #
    # This command retrieves a list of all projects.
    # You can also switch to a different project.
    def cmd_projects(self, *args):
        def usage():
            print("usage: projects [-h] [-l] [-s=project]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--list (-l)\tList all existing projects")
            print("\t--switch (-s)\tSwitch to the specified project")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hls:',
                                       ['help', 'list', 'switch='])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_list = False
        arg_switch = None

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-l', '--list'):
                arg_list = True
            elif opt in ('-s', '--switch'):
                arg_switch = value

        projects_path = os.path.join(os.getcwd(), 'projects')

        if not os.path.exists(projects_path):
            print_info("The projects directory does not exist yet")
            return

        if arg_list:
            print_info("Projects Available:")

            rows = []
            for project in os.listdir(projects_path):
                project_path = os.path.join(projects_path, project)
                if os.path.isdir(project_path):
                    current = ''
                    if __project__.name and project == __project__.name:
                        current = 'Yes'
                    rows.append([
                        project,
                        time.ctime(os.path.getctime(project_path)), current
                    ])

            print(
                table(header=['Project Name', 'Creation Time', 'Current'],
                      rows=rows))
            return
        elif arg_switch:
            if __sessions__.is_set():
                __sessions__.close()
                print_info("Closed opened session")

            __project__.open(arg_switch)
            print_info("Switched to project {0}".format(bold(arg_switch)))

            # Need to re-initialize the Database to open the new SQLite file.
            self.db = Database()
            return

        usage()

    ##
    # EXPORT
    #
    # This command will export the current session to file or zip.
    def cmd_export(self, *args):
        def usage():
            print("usage: export [-h] [-z] <path or archive name>")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--zip (-z)\tExport session in a zip archive")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hz', ['help', 'zip'])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_zip = False

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-z', '--zip'):
                arg_zip = True

        # This command requires a session to be opened.
        if not __sessions__.is_set():
            print_error("No session opened")
            return

        # Check for valid export path.
        if len(args) == 0:
            usage()
            return

        # TODO: having for one a folder and for the other a full
        # target path can be confusing. We should perhaps standardize this.

        # Abort if the specified path already exists.
        if os.path.isfile(argv[0]):
            print_error("File at path \"{0}\" already exists, abort".format(
                argv[0]))
            return

        # If the argument chosed so, archive the file when exporting it.
        # TODO: perhaps add an option to use a password for the archive
        # and default it to "infected".
        if arg_zip:
            try:
                with ZipFile(argv[0], 'w') as export_zip:
                    export_zip.write(__sessions__.current.file.path,
                                     arcname=__sessions__.current.file.name)
            except IOError as e:
                print_error("Unable to export file: {0}".format(e))
            else:
                print_info("File archived and exported to {0}".format(argv[0]))
        # Otherwise just dump it to the given directory.
        else:
            # XXX: Export file with the original file name.
            store_path = os.path.join(argv[0], __sessions__.current.file.name)

            try:
                shutil.copyfile(__sessions__.current.file.path, store_path)
            except IOError as e:
                print_error("Unable to export file: {0}".format(e))
            else:
                print_info("File exported to {0}".format(store_path))
示例#22
0
class Commands(object):

    output = []

    def __init__(self):
        # Open connection to the database.
        self.db = Database()

        # Map commands to their related functions.
        self.commands = dict(
            help=dict(obj=self.cmd_help, description="Show this help message"),
            open=dict(obj=self.cmd_open, description="Open a file"),
            new=dict(obj=self.cmd_new, description="Create new file"),
            close=dict(obj=self.cmd_close, description="Close the current session"),
            info=dict(obj=self.cmd_info, description="Show information on the opened file"),
            notes=dict(obj=self.cmd_notes, description="View, add and edit notes on the opened file"),
            clear=dict(obj=self.cmd_clear, description="Clear the console"),
            store=dict(obj=self.cmd_store, description="Store the opened file to the local repository"),
            delete=dict(obj=self.cmd_delete, description="Delete the opened file"),
            find=dict(obj=self.cmd_find, description="Find a file"),
            tags=dict(obj=self.cmd_tags, description="Modify tags of the opened file"),
            sessions=dict(obj=self.cmd_sessions, description="List or switch sessions"),
            stats=dict(obj=self.cmd_stats, description="Viper Collection Statistics"),
            projects=dict(obj=self.cmd_projects, description="List or switch existing projects"),
            export=dict(obj=self.cmd_export, description="Export the current session to file or zip"),
        )

    # Output Logging
    def log(self, event_type, event_data):
        self.output.append(dict(type=event_type, data=event_data))

    ##
    # CLEAR
    #
    # This command simply clears the shell.
    def cmd_clear(self, *args):
        os.system("clear")

    ##
    # HELP
    #
    # This command simply prints the help message.
    # It lists both embedded commands and loaded modules.
    def cmd_help(self, *args):
        self.log("info", "Commands")

        rows = []
        for command_name, command_item in self.commands.items():
            rows.append([command_name, command_item["description"]])

        rows.append(["exit, quit", "Exit Viper"])
        rows = sorted(rows, key=lambda entry: entry[0])

        self.log("table", dict(header=["Command", "Description"], rows=rows))
        self.log("info", "Modules")

        rows = []
        for module_name, module_item in __modules__.items():
            rows.append([module_name, module_item["description"]])

        rows = sorted(rows, key=lambda entry: entry[0])

        self.log("table", dict(header=["Command", "Description"], rows=rows))

    ##
    # NEW
    #
    # This command is used to create a new session on a new file,
    # useful for copy & paste of content like Email headers

    def cmd_new(self, *args):
        title = input("Enter a title for the new file: ")
        # Create a new temporary file.
        tmp = tempfile.NamedTemporaryFile(delete=False)
        # Open the temporary file with the default editor, or with nano.
        os.system('"${EDITOR:-nano}" ' + tmp.name)
        __sessions__.new(tmp.name)
        __sessions__.current.file.name = title
        print_info('New file with title "{0}" added to the current session'.format(bold(title)))

    ##
    # OPEN
    #
    # This command is used to open a session on a given file.
    # It either can be an external file path, or a SHA256 hash of a file which
    # has been previously imported and stored.
    # While the session is active, every operation and module executed will be
    # run against the file specified.
    def cmd_open(self, *args):
        parser = argparse.ArgumentParser(
            prog="open",
            description="Open a file",
            epilog="You can also specify a MD5 or SHA256 hash to a previously stored file in order to open a session on it.",
        )
        group = parser.add_mutually_exclusive_group()
        group.add_argument("-f", "--file", action="store_true", help="Target is a file")
        group.add_argument("-u", "--url", action="store_true", help="Target is a URL")
        group.add_argument(
            "-l", "--last", action="store_true", help="Target is the entry number from the last find command's results"
        )
        parser.add_argument("-t", "--tor", action="store_true", help="Download the file through Tor")
        parser.add_argument(
            "value",
            metavar="PATH, URL, HASH or ID",
            nargs="*",
            help="Target to open. Hash can be md5 or sha256. ID has to be from the last search.",
        )

        try:
            args = parser.parse_args(args)
        except:
            return

        target = " ".join(args.value)

        if not args.last and target is None:
            parser.print_usage()
            return

        # If it's a file path, open a session on it.
        if args.file:
            target = os.path.expanduser(target)

            if not os.path.exists(target) or not os.path.isfile(target):
                self.log("error", "File not found: {0}".format(target))
                return

            __sessions__.new(target)
        # If it's a URL, download it and open a session on the temporary file.
        elif args.url:
            data = download(url=target, tor=args.tor)

            if data:
                tmp = tempfile.NamedTemporaryFile(delete=False)
                tmp.write(data)
                tmp.close()

                __sessions__.new(tmp.name)
        # Try to open the specified file from the list of results from
        # the last find command.
        elif args.last:
            if __sessions__.find:
                count = 1
                for item in __sessions__.find:
                    if count == int(target):
                        __sessions__.new(get_sample_path(item.sha256))
                        break

                    count += 1
            else:
                self.log("warning", "You haven't performed a find yet")
        # Otherwise we assume it's an hash of an previously stored sample.
        else:
            target = target.strip().lower()

            if len(target) == 32:
                key = "md5"
            elif len(target) == 64:
                key = "sha256"
            else:
                parser.print_usage()
                return

            rows = self.db.find(key=key, value=target)

            if not rows:
                self.log("warning", "No file found with the given hash {0}".format(target))
                return

            path = get_sample_path(rows[0].sha256)
            if path:
                __sessions__.new(path)

    ##
    # CLOSE
    #
    # This command resets the open session.
    # After that, all handles to the opened file should be closed and the
    # shell should be restored to the default prompt.
    def cmd_close(self, *args):
        __sessions__.close()

    ##
    # INFO
    #
    # This command returns information on the open session. It returns details
    # on the file (e.g. hashes) and other information that might available from
    # the database.
    def cmd_info(self, *args):
        if __sessions__.is_set():
            self.log(
                "table",
                dict(
                    header=["Key", "Value"],
                    rows=[
                        ["Name", __sessions__.current.file.name],
                        ["Tags", __sessions__.current.file.tags],
                        ["Path", __sessions__.current.file.path],
                        ["Size", __sessions__.current.file.size],
                        ["Type", __sessions__.current.file.type],
                        ["Mime", __sessions__.current.file.mime],
                        ["MD5", __sessions__.current.file.md5],
                        ["SHA1", __sessions__.current.file.sha1],
                        ["SHA256", __sessions__.current.file.sha256],
                        ["SHA512", __sessions__.current.file.sha512],
                        ["SSdeep", __sessions__.current.file.ssdeep],
                        ["CRC32", __sessions__.current.file.crc32],
                    ],
                ),
            )

    ##
    # NOTES
    #
    # This command allows you to view, add, modify and delete notes associated
    # with the currently opened file.
    def cmd_notes(self, *args):
        parser = argparse.ArgumentParser(prog="notes", description="Show information on the opened file")
        group = parser.add_mutually_exclusive_group()
        group.add_argument("-l", "--list", action="store_true", help="List all notes available for the current file")
        group.add_argument("-a", "--add", action="store_true", help="Add a new note to the current file")
        group.add_argument("-v", "--view", metavar="NOTE ID", type=int, help="View the specified note")
        group.add_argument("-e", "--edit", metavar="NOTE ID", type=int, help="Edit an existing note")
        group.add_argument("-d", "--delete", metavar="NOTE ID", type=int, help="Delete an existing note")

        try:
            args = parser.parse_args(args)
        except:
            return

        if not __sessions__.is_set():
            self.log("error", "No session opened")
            return

        # check if the file is already stores, otherwise exit as no notes command will work if the file is not stored in the database
        malware = Database().find(key="sha256", value=__sessions__.current.file.sha256)
        if not malware:
            self.log("error", "The opened file doesn't appear to be in the database, have you stored it yet?")
            return

        if args.list:
            # Retrieve all notes for the currently opened file.

            notes = malware[0].note
            if not notes:
                self.log("info", "No notes available for this file yet")
                return

            # Populate table rows.
            rows = [[note.id, note.title] for note in notes]

            # Display list of existing notes.
            self.log("table", dict(header=["ID", "Title"], rows=rows))

        elif args.add:
            title = input("Enter a title for the new note: ")

            # Create a new temporary file.
            tmp = tempfile.NamedTemporaryFile(delete=False)
            # Open the temporary file with the default editor, or with nano.
            os.system('"${EDITOR:-nano}" ' + tmp.name)
            # Once the user is done editing, we need to read the content and
            # store it in the database.
            body = tmp.read()
            Database().add_note(__sessions__.current.file.sha256, title, body)
            # Finally, remove the temporary file.
            os.remove(tmp.name)

            self.log("info", 'New note with title "{0}" added to the current file'.format(bold(title)))

        elif args.view:
            # Retrieve note wth the specified ID and print it.
            note = Database().get_note(args.view)
            if note:
                self.log("info", bold("Title: ") + note.title)
                self.log("info", bold("Body:") + "\n" + note.body)
            else:
                self.log("info", "There is no note with ID {0}".format(args.view))

        elif args.edit:
            # Retrieve note with the specified ID.
            note = Database().get_note(args.edit)
            if note:
                # Create a new temporary file.
                tmp = tempfile.NamedTemporaryFile(delete=False)
                # Write the old body to the temporary file.
                tmp.write(note.body)
                tmp.close()
                # Open the old body with the text editor.
                os.system('"${EDITOR:-nano}" ' + tmp.name)
                # Read the new body from the temporary file.
                body = open(tmp.name, "r").read()
                # Update the note entry with the new body.
                Database().edit_note(args.edit, body)
                # Remove the temporary file.
                os.remove(tmp.name)

                self.log("info", "Updated note with ID {0}".format(args.edit))

        elif args.delete:
            # Delete the note with the specified ID.
            Database().delete_note(args.delete)
        else:
            parser.print_usage()

    ##
    # STORE
    #
    # This command stores the opened file in the local repository and tries
    # to store details in the database.
    def cmd_store(self, *args):
        parser = argparse.ArgumentParser(prog="store", description="Store the opened file to the local repository")
        parser.add_argument("-d", "--delete", action="store_true", help="Delete the original file")
        parser.add_argument("-f", "--folder", type=str, nargs="+", help="Specify a folder to import")
        parser.add_argument("-s", "--file-size", type=int, help="Specify a maximum file size")
        parser.add_argument("-y", "--file-type", type=str, help="Specify a file type pattern")
        parser.add_argument("-n", "--file-name", type=str, help="Specify a file name pattern")
        parser.add_argument("-t", "--tags", type=str, nargs="+", help="Specify a list of comma-separated tags")

        try:
            args = parser.parse_args(args)
        except:
            return

        if args.folder is not None:
            # Allows to have spaces in the path.
            args.folder = " ".join(args.folder)

        if args.tags is not None:
            # Remove the spaces in the list of tags
            args.tags = "".join(args.tags)

        def add_file(obj, tags=None):
            if get_sample_path(obj.sha256):
                self.log("warning", 'Skip, file "{0}" appears to be already stored'.format(obj.name))
                return False

            # Try to store file object into database.
            status = self.db.add(obj=obj, tags=tags)
            if status:
                # If succeeds, store also in the local repository.
                # If something fails in the database (for example unicode strings)
                # we don't want to have the binary lying in the repository with no
                # associated database record.
                new_path = store_sample(obj)
                self.log("success", 'Stored file "{0}" to {1}'.format(obj.name, new_path))
            else:
                return False

            # Delete the file if requested to do so.
            if args.delete:
                try:
                    os.unlink(obj.path)
                except Exception as e:
                    self.log("warning", "Failed deleting file: {0}".format(e))

            return True

        # If the user specified the --folder flag, we walk recursively and try
        # to add all contained files to the local repository.
        # This is note going to open a new session.
        # TODO: perhaps disable or make recursion optional?
        if args.folder is not None:
            # Check if the specified folder is valid.
            if os.path.isdir(args.folder):
                # Walk through the folder and subfolders.
                for dir_name, dir_names, file_names in walk(args.folder):
                    # Add each collected file.
                    for file_name in file_names:
                        file_path = os.path.join(dir_name, file_name)

                        if not os.path.exists(file_path):
                            continue
                        # Check if file is not zero.
                        if not os.path.getsize(file_path) > 0:
                            continue

                        # Check if the file name matches the provided pattern.
                        if args.file_name:
                            if not fnmatch.fnmatch(file_name, args.file_name):
                                # self.log('warning', "Skip, file \"{0}\" doesn't match the file name pattern".format(file_path))
                                continue

                        # Check if the file type matches the provided pattern.
                        if args.file_type:
                            if args.file_type not in File(file_path).type:
                                # self.log('warning', "Skip, file \"{0}\" doesn't match the file type".format(file_path))
                                continue

                        # Check if file exceeds maximum size limit.
                        if args.file_size:
                            # Obtain file size.
                            if os.path.getsize(file_path) > args.file_size:
                                self.log("warning", 'Skip, file "{0}" is too big'.format(file_path))
                                continue

                        file_obj = File(file_path)

                        # Add file.
                        add_file(file_obj, args.tags)
            else:
                self.log("error", "You specified an invalid folder: {0}".format(args.folder))
        # Otherwise we try to store the currently opened file, if there is any.
        else:
            if __sessions__.is_set():
                if __sessions__.current.file.size == 0:
                    self.log("warning", 'Skip, file "{0}" appears to be empty'.format(__sessions__.current.file.name))
                    return False

                # Add file.
                if add_file(__sessions__.current.file, args.tags):
                    # Open session to the new file.
                    self.cmd_open(*[__sessions__.current.file.sha256])
            else:
                self.log("error", "No session opened")

    ##
    # DELETE
    #
    # This commands deletes the currenlty opened file (only if it's stored in
    # the local repository) and removes the details from the database
    def cmd_delete(self, *args):
        if __sessions__.is_set():
            while True:
                choice = input("Are you sure you want to delete this binary? Can't be reverted! [y/n] ")
                if choice == "y":
                    break
                elif choice == "n":
                    return

            rows = self.db.find("sha256", __sessions__.current.file.sha256)
            if rows:
                malware_id = rows[0].id
                if self.db.delete_file(malware_id):
                    self.log("success", "File deleted")
                else:
                    self.log("error", "Unable to delete file")

            os.remove(__sessions__.current.file.path)
            __sessions__.close()
        else:
            self.log("error", "No session opened")

    ##
    # FIND
    #
    # This command is used to search for files in the database.
    def cmd_find(self, *args):
        parser = argparse.ArgumentParser(prog="find", description="Find a file")
        group = parser.add_mutually_exclusive_group()
        group.add_argument("-t", "--tags", action="store_true", help="List available tags and quit")
        group.add_argument(
            "type",
            nargs="?",
            choices=["all", "latest", "name", "type", "mime", "md5", "sha256", "tag", "note"],
            help="Where to search.",
        )
        parser.add_argument("value", nargs="?", help="String to search.")
        try:
            args = parser.parse_args(args)
        except:
            return

        # One of the most useful search terms is by tag. With the --tags
        # argument we first retrieve a list of existing tags and the count
        # of files associated with each of them.
        if args.tags:
            # Retrieve list of tags.
            tags = self.db.list_tags()

            if tags:
                rows = []
                # For each tag, retrieve the count of files associated with it.
                for tag in tags:
                    count = len(self.db.find("tag", tag.tag))
                    rows.append([tag.tag, count])

                # Generate the table with the results.
                header = ["Tag", "# Entries"]
                rows.sort(key=lambda x: x[1], reverse=True)
                self.log("table", dict(header=header, rows=rows))
            else:
                self.log("warning", "No tags available")

            return

        # At this point, if there are no search terms specified, return.
        if args.type is None:
            parser.print_usage()
            return

        key = args.type
        if key != "all" and key != "latest":
            try:
                # The second argument is the search value.
                value = args.value
            except IndexError:
                self.log("error", "You need to include a search term.")
                return
        else:
            value = None

        # Search all the files matching the given parameters.
        items = self.db.find(key, value)
        if not items:
            return

        # Populate the list of search results.
        rows = []
        count = 1
        for item in items:
            tag = ", ".join([t.tag for t in item.tag if t.tag])
            row = [count, item.name, item.mime, item.md5, tag]
            if key == "latest":
                row.append(item.created_at)

            rows.append(row)
            count += 1

        # Update find results in current session.
        __sessions__.find = items

        # Generate a table with the results.
        header = ["#", "Name", "Mime", "MD5", "Tags"]
        if key == "latest":
            header.append("Created At")

        self.log("table", dict(header=header, rows=rows))

    ##
    # TAGS
    #
    # This command is used to modify the tags of the opened file.
    def cmd_tags(self, *args):
        parser = argparse.ArgumentParser(prog="tags", description="Modify tags of the opened file")
        parser.add_argument("-a", "--add", metavar="TAG", help="Add tags to the opened file (comma separated)")
        parser.add_argument("-d", "--delete", metavar="TAG", help="Delete a tag from the opened file")
        try:
            args = parser.parse_args(args)
        except:
            return

        # This command requires a session to be opened.
        if not __sessions__.is_set():
            self.log("error", "No session opened")
            parser.print_usage()
            return

        # If no arguments are specified, there's not much to do.
        # However, it could make sense to also retrieve a list of existing
        # tags from this command, and not just from the "find" command alone.
        if args.add is None and args.delete is None:
            parser.print_usage()
            return

        # TODO: handle situation where addition or deletion of a tag fail.

        db = Database()
        if not db.find(key="sha256", value=__sessions__.current.file.sha256):
            self.log(
                "error",
                "The opened file is not stored in the database. " "If you want to add it use the `store` command.",
            )
            return

        if args.add:
            # Add specified tags to the database's entry belonging to
            # the opened file.
            db.add_tags(__sessions__.current.file.sha256, args.add)
            self.log("info", "Tags added to the currently opened file")

            # We refresh the opened session to update the attributes.
            # Namely, the list of tags returned by the 'info' command
            # needs to be re-generated, or it wouldn't show the new tags
            # until the existing session is closed a new one is opened.
            self.log("info", "Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

        if args.delete:
            # Delete the tag from the database.
            db.delete_tag(args.delete, __sessions__.current.file.sha256)
            # Refresh the session so that the attributes of the file are
            # updated.
            self.log("info", "Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

    ###
    # SESSION
    #
    # This command is used to list and switch across all the opened sessions.
    def cmd_sessions(self, *args):
        parser = argparse.ArgumentParser(prog="sessions", description="Open a file", epilog="List or switch sessions")
        group = parser.add_mutually_exclusive_group()
        group.add_argument("-l", "--list", action="store_true", help="List all existing sessions")
        group.add_argument("-s", "--switch", type=int, help="Switch to the specified session")

        try:
            args = parser.parse_args(args)
        except:
            return

        if args.list:
            if not __sessions__.sessions:
                self.log("info", "There are no opened sessions")
                return

            rows = []
            for session in __sessions__.sessions:
                current = ""
                if session == __sessions__.current:
                    current = "Yes"

                rows.append([session.id, session.file.name, session.file.md5, session.created_at, current])

            self.log("info", "Opened Sessions:")
            self.log("table", dict(header=["#", "Name", "MD5", "Created At", "Current"], rows=rows))
        elif args.switch:
            for session in __sessions__.sessions:
                if args.switch == session.id:
                    __sessions__.switch(session)
                    return

            self.log("warning", "The specified session ID doesn't seem to exist")
        else:
            parser.print_usage()

    ##
    # PROJECTS
    #
    # This command retrieves a list of all projects.
    # You can also switch to a different project.
    def cmd_projects(self, *args):
        parser = argparse.ArgumentParser(
            prog="projects", description="Open a file", epilog="List or switch existing projects"
        )
        group = parser.add_mutually_exclusive_group()
        group.add_argument("-l", "--list", action="store_true", help="List all existing projects")
        group.add_argument("-s", "--switch", metavar="PROJECT NAME", help="Switch to the specified project")

        try:
            args = parser.parse_args(args)
        except:
            return

        projects_path = os.path.join(os.getcwd(), "projects")

        if not os.path.exists(projects_path):
            self.log("info", "The projects directory does not exist yet")
            return

        if args.list:
            self.log("info", "Projects Available:")

            rows = []
            for project in os.listdir(projects_path):
                project_path = os.path.join(projects_path, project)
                if os.path.isdir(project_path):
                    current = ""
                    if __project__.name and project == __project__.name:
                        current = "Yes"
                    rows.append([project, time.ctime(os.path.getctime(project_path)), current])

            self.log("table", dict(header=["Project Name", "Creation Time", "Current"], rows=rows))
        elif args.switch:
            if __sessions__.is_set():
                __sessions__.close()
                self.log("info", "Closed opened session")

            __project__.open(args.switch)
            self.log("info", "Switched to project {0}".format(bold(args.switch)))

            # Need to re-initialize the Database to open the new SQLite file.
            self.db = Database()
        else:
            self.log("info", parser.print_usage())

    ##
    # EXPORT
    #
    # This command will export the current session to file or zip.
    def cmd_export(self, *args):
        parser = argparse.ArgumentParser(prog="export", description="Export the current session to file or zip")
        parser.add_argument("-z", "--zip", action="store_true", help="Export session in a zip archive")
        parser.add_argument("value", help="path or archive name")

        try:
            args = parser.parse_args(args)
        except:
            return

        # This command requires a session to be opened.
        if not __sessions__.is_set():
            self.log("error", "No session opened")
            parser.print_usage()
            return

        # Check for valid export path.
        if args.value is None:
            parser.print_usage()
            return

        # TODO: having for one a folder and for the other a full
        # target path can be confusing. We should perhaps standardize this.

        # Abort if the specified path already exists.
        if os.path.isfile(args.value):
            self.log("error", 'File at path "{0}" already exists, abort'.format(args.value))
            return

        # If the argument chosed so, archive the file when exporting it.
        # TODO: perhaps add an option to use a password for the archive
        # and default it to "infected".
        if args.zip:
            try:
                with ZipFile(args.value, "w") as export_zip:
                    export_zip.write(__sessions__.current.file.path, arcname=__sessions__.current.file.name)
            except IOError as e:
                self.log("error", "Unable to export file: {0}".format(e))
            else:
                self.log("info", "File archived and exported to {0}".format(args.value))
        # Otherwise just dump it to the given directory.
        else:
            # XXX: Export file with the original file name.
            store_path = os.path.join(args.value, __sessions__.current.file.name)

            try:
                shutil.copyfile(__sessions__.current.file.path, store_path)
            except IOError as e:
                self.log("error", "Unable to export file: {0}".format(e))
            else:
                self.log("info", "File exported to {0}".format(store_path))

    ##
    # Stats
    #
    # This command allows you to generate basic statistics for the stored files.
    def cmd_stats(self, *args):
        parser = argparse.ArgumentParser(prog="stats", description="Display Database File Statistics")
        parser.add_argument("-t", "--top", type=int, help="Top x Items")

        try:
            args = parser.parse_args(args)
        except:
            return

        arg_top = args.top
        db = Database()

        # Set all Counters Dict
        extension_dict = defaultdict(int)
        mime_dict = defaultdict(int)
        tags_dict = defaultdict(int)
        size_list = []

        # Find all
        items = self.db.find("all")

        if len(items) < 1:
            self.log("info", "No items in database to generate stats")
            return

        # Sort in to stats
        for item in items:
            if "." in item.name:
                ext = item.name.split(".")
                extension_dict[ext[-1]] += 1
            mime_dict[item.mime] += 1
            size_list.append(item.size)
            for t in item.tag:
                if t.tag:
                    tags_dict[t.tag] += 1

        avg_size = sum(size_list) / len(size_list)
        all_stats = {
            "Total": len(items),
            "File Extension": extension_dict,
            "Mime": mime_dict,
            "Tags": tags_dict,
            "Avg Size": avg_size,
            "Largest": max(size_list),
            "Smallest": min(size_list),
        }

        # Counter for top x
        if arg_top:
            counter = arg_top
            prefix = "Top {0} ".format(counter)
        else:
            counter = len(items)
            prefix = ""

        # Project Stats Last as i have it iterate them all

        # Print all the results

        self.log("info", "Projects")
        self.log("table", dict(header=["Name", "Count"], rows=[["Main", len(items)], ["Next", "10"]]))

        # For Current Project
        self.log("info", "Current Project")

        # Extension
        self.log("info", "{0}Extensions".format(prefix))
        header = ["Ext", "Count"]
        rows = []

        for k in sorted(extension_dict, key=extension_dict.get, reverse=True)[:counter]:
            rows.append([k, extension_dict[k]])
        self.log("table", dict(header=header, rows=rows))

        # Mimes
        self.log("info", "{0}Mime Types".format(prefix))
        header = ["Mime", "Count"]
        rows = []
        for k in sorted(mime_dict, key=mime_dict.get, reverse=True)[:counter]:
            rows.append([k, mime_dict[k]])
        self.log("table", dict(header=header, rows=rows))

        # Tags
        self.log("info", "{0}Tags".format(prefix))
        header = ["Tag", "Count"]
        rows = []
        for k in sorted(tags_dict, key=tags_dict.get, reverse=True)[:counter]:
            rows.append([k, tags_dict[k]])
        self.log("table", dict(header=header, rows=rows))

        # Size
        self.log("info", "Size Stats")
        self.log("item", "Largest  {0}".format(convert_size(max(size_list))))
        self.log("item", "Smallest  {0}".format(convert_size(min(size_list))))
        self.log("item", "Average  {0}".format(convert_size(avg_size)))
示例#23
0
    def test_add(self, capsys, filename, name):
        f = File(os.path.join(FIXTURE_DIR, filename))

        instance = Database()
        ret = instance.add(f)
        assert ret is True
示例#24
0
class Commands(object):

    def __init__(self):
        # Open connection to the database.
        self.db = Database()

        # Map commands to their related functions.
        self.commands = dict(
            help=dict(obj=self.cmd_help, description="Show this help message"),
            open=dict(obj=self.cmd_open, description="Open a file"),
            close=dict(obj=self.cmd_close, description="Close the current session"),
            info=dict(obj=self.cmd_info, description="Show information on the opened file"),
            clear=dict(obj=self.cmd_clear, description="Clear the console"),
            store=dict(obj=self.cmd_store, description="Store the opened file to the local repository"),
            delete=dict(obj=self.cmd_delete, description="Delete the opened file"),
            find=dict(obj=self.cmd_find, description="Find a file"),
        )

    ##
    # CLEAR
    #
    # This command simply clears the shell.
    def cmd_clear(self, *args):
        os.system('clear')

    ##
    # HELP
    #
    # This command simply prints the help message.
    # It lists both embedded commands and loaded modules.
    def cmd_help(self, *args):
        print(bold("Commands:"))

        rows = []
        for command_name, command_item in self.commands.items():
            rows.append([command_name, command_item['description']])

        print(table(['Command', 'Description'], rows))       
        print("")
        print(bold("Modules:"))

        rows = []
        for module_name, module_item in __modules__.items():
            rows.append([module_name, module_item['description']])

        print(table(['Command', 'Description'], rows))

    ##
    # OPEN
    #
    # This command is used to open a session on a given file.
    # It either can be an external file path, or a SHA256 hash of a file which
    # has been previously imported and stored.
    # While the session is active, every operation and module executed will be
    # run against the file specified.
    def cmd_open(self, *args):
        def usage():
            print("usage: open [-h] [-f] [-u] [-t] <target>")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--file (-f)\tThe target is a file")
            print("\t--url (-u)\tThe target is a URL")
            print("\t--tor (-t)\tDownload the file through Tor")
            print("")
            print("You can also specify a SHA256 hash to a previously stored")
            print("file in order to open a session on it.")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hfut', ['help', 'file', 'url', 'tor'])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        is_file = False
        is_url = False
        use_tor = False

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-f', '--file'):
                is_file = True
            elif opt in ('-u', '--url'):
                is_url = True
            elif opt in ('-t', '--tor'):
                use_tor = True

        if len(argv) == 0:
            usage()
            return
        else:
            target = argv[0]

        # If it's a file path, open a session on it.
        if is_file:
            target = os.path.expanduser(target)

            if not os.path.exists(target) or not os.path.isfile(target):
                print_error("File not found")
                return

            __session__.set(target)
        # If it's a URL, download it and open a session on the temporary
        # file.
        elif is_url:
            data = download(url=target, tor=use_tor)

            if data:
                tmp = tempfile.NamedTemporaryFile(delete=False)
                tmp.write(data)
                tmp.close()

                __session__.set(tmp.name)
        # Otherwise we assume it's an hash of an previously stored sample.
        else:
            target = argv[0].strip().lower()
            path = get_sample_path(target)
            if path:
                __session__.set(path)

    ##
    # CLOSE
    #
    # This command resets the open session.
    # After that, all handles to the opened file should be closed and the
    # shell should be restored to the default prompt.
    def cmd_close(self, *args):
        __session__.clear()

    ##
    # INFO
    #
    # This command returns information on the open session. It returns details
    # on the file (e.g. hashes) and other information that might available from
    # the database.
    def cmd_info(self, *args):
        if __session__.is_set():
            print(table(
                ['Key', 'Value'],
                [
                    ('Name', __session__.file.name),
                    ('Tags', __session__.file.tags),
                    ('Path', __session__.file.path),
                    ('Size', __session__.file.size),
                    ('Type', __session__.file.type),
                    ('MD5', __session__.file.md5),
                    ('SHA1', __session__.file.sha1),
                    ('SHA256', __session__.file.sha256),
                    ('SHA512', __session__.file.sha512),
                    ('SSdeep', __session__.file.ssdeep),
                    ('CRC32', __session__.file.crc32)
                ]
            ))

    ##
    # STORE
    #
    # This command stores the opened file in the local repository and tries
    # to store details in the database.
    def cmd_store(self, *args):
        def usage():
            print("usage: store [-h] [-d] [-f <path>] [-t]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--delete (-d)\tDelete the original file")
            print("\t--folder (-f)\tSpecify a folder to import")
            print("\t--tags (-t)\tSpecify a list of comma-separated tags")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hdf:t:', ['help', 'delete', 'folder=', 'tags='])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        do_delete = False
        folder = False
        tags = None

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-d', '--delete'):
                do_delete = True
            elif opt in ('-f', '--folder'):
                folder = value
            elif opt in ('-t', '--tags'):
                tags = value

        def add_file(obj, tags=None):
            # Store file to the local repository.
            new_path = store_sample(obj)
            if new_path:
                # Add file to the database.
                status = self.db.add(obj=obj, tags=tags)
                print_success("Stored to: {0}".format(new_path))

            # Delete the file if requested to do so.
            if do_delete:
                try:
                    os.unlink(obj.path)
                except Exception as e:
                    print_warning("Failed deleting file: {0}".format(e))

        # If the user specified the --folder flag, we walk recursively and try
        # to add all contained files to the local repository.
        # This is note going to open a new session.
        # TODO: perhaps disable or make recursion optional?
        if folder:
            # Check if the specified folder is valid.
            if os.path.isdir(folder):
                # Walk through the folder and subfolders.
                for dir_name, dir_names, file_names in os.walk(folder):
                    # Add each collected file.
                    for file_name in file_names:
                        file_path = os.path.join(dir_name, file_name)
                        file_obj = File(file_path)

                        # Add file.
                        add_file(file_obj, tags)
            else:
                print_error("You specified an invalid folder: {0}".format(folder))
        # Otherwise we try to store the currently opened file, if there is any.
        else:
            if __session__.is_set():
                # Add file.
                add_file(__session__.file, tags)

                # Open session to the new file.
                self.cmd_open(*[__session__.file.sha256])
            else:
                print_error("No session opened")

    ##
    # DELETE
    #
    # This commands deletes the currenlty opened file (only if it's stored in
    # the local repository) and removes the details from the database
    def cmd_delete(self, *args):
        if __session__.is_set():
            while True:
                choice = raw_input("Are you sure you want to delete this binary? Can't be reverted! [y/n] ")
                if choice == 'y':
                    break
                elif choice == 'n':
                    return

            rows = self.db.find('sha256', __session__.file.sha256)
            if rows:
                malware_id = rows[0].id
                if self.db.delete(malware_id):
                    print_success("File deleted")
                else:
                    print_error("Unable to delete file")

            os.remove(get_sample_path(__session__.file.sha256))
            __session__.clear()

    ##
    # FIND
    #
    # This command is used to search for files in the database.
    def cmd_find(self, *args):
        if len(args) == 0:
            print_error("Invalid search term")
            return

        key = args[0]
        try:
            value = args[1]
        except IndexError:
            value = None

        items = self.db.find(key, value)
        if not items:
            return

        rows = []
        for item in items:
            rows.append([item.name, item.type, item.sha256])

        print(table(['Name', 'Type',  'SHA256'], rows))
示例#25
0
文件: commands.py 项目: creturn/viper
class Commands(object):

    def __init__(self):
        # Open connection to the database.
        self.db = Database()

        # Map commands to their related functions.
        self.commands = dict(
            help=dict(obj=self.cmd_help, description="Show this help message"),
            open=dict(obj=self.cmd_open, description="Open a file"),
            close=dict(obj=self.cmd_close, description="Close the current session"),
            info=dict(obj=self.cmd_info, description="Show information on the opened file"),
            notes=dict(obj=self.cmd_notes, description="View, add and edit notes on the opened file"),
            clear=dict(obj=self.cmd_clear, description="Clear the console"),
            store=dict(obj=self.cmd_store, description="Store the opened file to the local repository"),
            delete=dict(obj=self.cmd_delete, description="Delete the opened file"),
            find=dict(obj=self.cmd_find, description="Find a file"),
            tags=dict(obj=self.cmd_tags, description="Modify tags of the opened file"),
            sessions=dict(obj=self.cmd_sessions, description="List or switch sessions"),
            projects=dict(obj=self.cmd_projects, description="List or switch existing projects"),
            export=dict(obj=self.cmd_export, description="Export the current session to file or zip"),
        )

    ##
    # CLEAR
    #
    # This command simply clears the shell.
    def cmd_clear(self, *args):
        os.system('clear')

    ##
    # HELP
    #
    # This command simply prints the help message.
    # It lists both embedded commands and loaded modules.
    def cmd_help(self, *args):
        print(bold("Commands:"))

        rows = []
        for command_name, command_item in self.commands.items():
            rows.append([command_name, command_item['description']])

        rows = sorted(rows, key=lambda entry: entry[0])

        print(table(['Command', 'Description'], rows))
        print("")
        print(bold("Modules:"))

        rows = []
        for module_name, module_item in __modules__.items():
            rows.append([module_name, module_item['description']])

        rows = sorted(rows, key=lambda entry: entry[0])

        print(table(['Command', 'Description'], rows))

    ##
    # OPEN
    #
    # This command is used to open a session on a given file.
    # It either can be an external file path, or a SHA256 hash of a file which
    # has been previously imported and stored.
    # While the session is active, every operation and module executed will be
    # run against the file specified.
    def cmd_open(self, *args):
        def usage():
            print("usage: open [-h] [-f] [-u] [-l] [-t] <target|md5|sha256>")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--file (-f)\tThe target is a file")
            print("\t--url (-u)\tThe target is a URL")
            print("\t--last (-l)\tThe target is the entry number from the last find command's results")
            print("\t--tor (-t)\tDownload the file through Tor")
            print("")
            print("You can also specify a MD5 or SHA256 hash to a previously stored")
            print("file in order to open a session on it.")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hfult', ['help', 'file', 'url', 'last', 'tor'])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_is_file = False
        arg_is_url = False
        arg_last = False
        arg_use_tor = False

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-f', '--file'):
                arg_is_file = True
            elif opt in ('-u', '--url'):
                arg_is_url = True
            elif opt in ('-l', '--last'):
                arg_last = True
            elif opt in ('-t', '--tor'):
                arg_use_tor = True

        if len(argv) == 0:
            usage()
            return
        else:
            target = argv[0]

        # If it's a file path, open a session on it.
        if arg_is_file:
            target = os.path.expanduser(target)

            # This is kind of hacky. It checks if there are additional arguments
            # to the open command, if there is I assume that it's the continuation
            # of a filename with spaces. I then concatenate them.
            # TODO: improve this.
            if len(argv) > 1:
                for arg in argv[1:]:
                    target += ' ' + arg

            if not os.path.exists(target) or not os.path.isfile(target):
                print_error("File not found: {0}".format(target))
                return

            __sessions__.new(target)
        # If it's a URL, download it and open a session on the temporary
        # file.
        elif arg_is_url:
            data = download(url=target, tor=arg_use_tor)

            if data:
                tmp = tempfile.NamedTemporaryFile(delete=False)
                tmp.write(data)
                tmp.close()

                __sessions__.new(tmp.name)
        # Try to open the specified file from the list of results from
        # the last find command.
        elif arg_last:
            if __sessions__.find:
                count = 1
                for item in __sessions__.find:
                    if count == int(target):
                        __sessions__.new(get_sample_path(item.sha256))
                        break

                    count += 1
            else:
                print_warning("You haven't performed a find yet")
        # Otherwise we assume it's an hash of an previously stored sample.
        else:
            target = argv[0].strip().lower()

            if len(target) == 32:
                key = 'md5'
            elif len(target) == 64:
                key = 'sha256'
            else:
                usage()
                return

            rows = self.db.find(key=key, value=target)

            if not rows:
                print_warning("No file found with the given hash {0}".format(target))
                return

            path = get_sample_path(rows[0].sha256)
            if path:
                __sessions__.new(path)

    ##
    # CLOSE
    #
    # This command resets the open session.
    # After that, all handles to the opened file should be closed and the
    # shell should be restored to the default prompt.
    def cmd_close(self, *args):
        __sessions__.close()

    ##
    # INFO
    #
    # This command returns information on the open session. It returns details
    # on the file (e.g. hashes) and other information that might available from
    # the database.
    def cmd_info(self, *args):
        if __sessions__.is_set():
            print(table(
                ['Key', 'Value'],
                [
                    ('Name', __sessions__.current.file.name),
                    ('Tags', __sessions__.current.file.tags),
                    ('Path', __sessions__.current.file.path),
                    ('Size', __sessions__.current.file.size),
                    ('Type', __sessions__.current.file.type),
                    ('Mime', __sessions__.current.file.mime),
                    ('MD5', __sessions__.current.file.md5),
                    ('SHA1', __sessions__.current.file.sha1),
                    ('SHA256', __sessions__.current.file.sha256),
                    ('SHA512', __sessions__.current.file.sha512),
                    ('SSdeep', __sessions__.current.file.ssdeep),
                    ('CRC32', __sessions__.current.file.crc32)
                ]
            ))

    ##
    # NOTES
    #
    # This command allows you to view, add, modify and delete notes associated
    # with the currently opened file.
    def cmd_notes(self, *args):
        def usage():
            print("usage: notes [-h] [-l] [-a] [-e <note id>] [-d <note id>]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--list (-l)\tList all notes available for the current file")
            print("\t--add (-a)\tAdd a new note to the current file")
            print("\t--view (-v)\tView the specified note")
            print("\t--edit (-e)\tEdit an existing note")
            print("\t--delete (-d)\tDelete an existing note")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hlav:e:d:', ['help', 'list', 'add', 'view=', 'edit=', 'delete='])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_list = False
        arg_add = False
        arg_view = None
        arg_edit = None
        arg_delete = None

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-l', '--list'):
                arg_list = True
            elif opt in ('-a', '--add'):
                arg_add = True
            elif opt in ('-v', '--view'):
                arg_view = value
            elif opt in ('-e', '--edit'):
                arg_edit = value
            elif opt in ('-d', '--delete'):
                arg_delete = value

        if not __sessions__.is_set():
            print_error("No session opened")
            return

        if arg_list:
            # Retrieve all notes for the currently opened file.
            malware = Database().find(key='sha256', value=__sessions__.current.file.sha256)
            if not malware:
                print_error("The opened file doesn't appear to be in the database, have you stored it yet?")
                return

            notes = malware[0].note
            if not notes:
                print_info("No notes available for this file yet")
                return

            # Populate table rows.
            rows = []
            for note in notes:
                rows.append([note.id, note.title])

            # Display list of existing notes.
            print(table(header=['ID', 'Title'], rows=rows))

        elif arg_add:
            title = raw_input("Enter a title for the new note: ")

            # Create a new temporary file.
            tmp = tempfile.NamedTemporaryFile(delete=False)
            # Open the temporary file with the default editor, or with nano.
            os.system('"${EDITOR:-nano}" ' + tmp.name)
            # Once the user is done editing, we need to read the content and
            # store it in the database.
            body = tmp.read()
            Database().add_note(__sessions__.current.file.sha256, title, body)
            # Finally, remove the temporary file.
            os.remove(tmp.name)

            print_info("New note with title \"{0}\" added to the current file".format(bold(title)))

        elif arg_view:
            # Retrieve note wth the specified ID and print it.
            note = Database().get_note(arg_view)
            if note:
                print_info(bold('Title: ') + note.title)
                print_info(bold('Body:'))
                print(note.body)
            else:
                print_info("There is no note with ID {0}".format(arg_view))


        elif arg_edit:
            # Retrieve note with the specified ID.
            note = Database().get_note(arg_edit)
            if note:
                # Create a new temporary file.
                tmp = tempfile.NamedTemporaryFile(delete=False)
                # Write the old body to the temporary file.
                tmp.write(note.body)
                tmp.close()
                # Open the old body with the text editor.
                os.system('"${EDITOR:-nano}" ' + tmp.name)
                # Read the new body from the temporary file.
                body = open(tmp.name, 'r').read()
                # Update the note entry with the new body.
                Database().edit_note(arg_edit, body)
                # Remove the temporary file.
                os.remove(tmp.name)

                print_info("Updated note with ID {0}".format(arg_edit))

        elif arg_delete:
            # Delete the note with the specified ID.
            Database().delete_note(arg_delete)
        else:
            usage()

    ##
    # STORE
    #
    # This command stores the opened file in the local repository and tries
    # to store details in the database.
    def cmd_store(self, *args):
        def usage():
            print("usage: store [-h] [-d] [-f <path>] [-s <size>] [-y <type>] [-n <name>] [-t]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--delete (-d)\tDelete the original file")
            print("\t--folder (-f)\tSpecify a folder to import")
            print("\t--file-size (-s)\tSpecify a maximum file size")
            print("\t--file-type (-y)\tSpecify a file type pattern")
            print("\t--file-name (-n)\tSpecify a file name pattern")
            print("\t--tags (-t)\tSpecify a list of comma-separated tags")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hdf:s:y:n:t:', ['help', 'delete', 'folder=', 'file-size=', 'file-type=', 'file-name=', 'tags='])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_delete = False
        arg_folder = False
        arg_file_size = None
        arg_file_type = None
        arg_file_name = None
        arg_tags = None

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-d', '--delete'):
                arg_delete = True
            elif opt in ('-f', '--folder'):
                arg_folder = value
            elif opt in ('-s', '--file-size'):
                arg_file_size = value
            elif opt in ('-y', '--file-type'):
                arg_file_type = value
            elif opt in ('-n', '--file-name'):
                arg_file_name = value
            elif opt in ('-t', '--tags'):
                arg_tags = value

        def add_file(obj, tags=None):
            if get_sample_path(obj.sha256):
                print_warning("Skip, file \"{0}\" appears to be already stored".format(obj.name))
                return False

            # Try to store file object into database.
            status = self.db.add(obj=obj, tags=tags)
            if status:
                # If succeeds, store also in the local repository.
                # If something fails in the database (for example unicode strings)
                # we don't want to have the binary lying in the repository with no
                # associated database record.
                new_path = store_sample(obj)
                print_success("Stored file \"{0}\" to {1}".format(obj.name, new_path))
            else:
                return False

            # Delete the file if requested to do so.
            if arg_delete:
                try:
                    os.unlink(obj.path)
                except Exception as e:
                    print_warning("Failed deleting file: {0}".format(e))

            return True

        # If the user specified the --folder flag, we walk recursively and try
        # to add all contained files to the local repository.
        # This is note going to open a new session.
        # TODO: perhaps disable or make recursion optional?
        if arg_folder:
            # Check if the specified folder is valid.
            if os.path.isdir(arg_folder):
                # Walk through the folder and subfolders.
                for dir_name, dir_names, file_names in os.walk(arg_folder):
                    # Add each collected file.
                    for file_name in file_names:
                        file_path = os.path.join(dir_name, file_name)

                        if not os.path.exists(file_path):
                            continue
                        # Check if file is not zero.
                        if not os.path.getsize(file_path) > 0:
                            continue

                        # Check if the file name matches the provided pattern.
                        if arg_file_name:
                            if not fnmatch.fnmatch(file_name, arg_file_name):
                                #print_warning("Skip, file \"{0}\" doesn't match the file name pattern".format(file_path))
                                continue

                        # Check if the file type matches the provided pattern.
                        if arg_file_type:
                            if arg_file_type not in File(file_path).type:
                                #print_warning("Skip, file \"{0}\" doesn't match the file type".format(file_path))
                                continue

                        # Check if file exceeds maximum size limit.
                        if arg_file_size:
                            # Obtain file size.
                            if os.path.getsize(file_path) > arg_file_size:
                                print_warning("Skip, file \"{0}\" is too big".format(file_path))
                                continue

                        file_obj = File(file_path)

                        # Add file.
                        add_file(file_obj, arg_tags)
            else:
                print_error("You specified an invalid folder: {0}".format(arg_folder))
        # Otherwise we try to store the currently opened file, if there is any.
        else:
            if __sessions__.is_set():
                if __sessions__.current.file.size == 0:
                    print_warning("Skip, file \"{0}\" appears to be empty".format(__sessions__.current.file.name))
                    return False

                # Add file.
                if add_file(__sessions__.current.file, arg_tags):
                    # Open session to the new file.
                    self.cmd_open(*[__sessions__.current.file.sha256])
            else:
                print_error("No session opened")

    ##
    # DELETE
    #
    # This commands deletes the currenlty opened file (only if it's stored in
    # the local repository) and removes the details from the database
    def cmd_delete(self, *args):
        if __sessions__.is_set():
            while True:
                choice = raw_input("Are you sure you want to delete this binary? Can't be reverted! [y/n] ")
                if choice == 'y':
                    break
                elif choice == 'n':
                    return

            rows = self.db.find('sha256', __sessions__.current.file.sha256)
            if rows:
                malware_id = rows[0].id
                if self.db.delete(malware_id):
                    print_success("File deleted")
                else:
                    print_error("Unable to delete file")

            os.remove(__sessions__.current.file.path)
            __sessions__.close()
        else:
            print_error("No session opened")

    ##
    # FIND
    #
    # This command is used to search for files in the database.
    def cmd_find(self, *args):
        def usage():
            print("usage: find [-h] [-t] <all|latest|name|type|mime|md5|sha256|tag|note> <value>")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--tags (-t)\tList tags")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'ht', ['help', 'tags'])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_list_tags = False

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-t', '--tags'):
                arg_list_tags = True

        # One of the most useful search terms is by tag. With the --tags
        # argument we first retrieve a list of existing tags and the count
        # of files associated with each of them.
        if arg_list_tags:
            # Retrieve list of tags.
            tags = self.db.list_tags()

            if tags:
                rows = []
                # For each tag, retrieve the count of files associated with it.
                for tag in tags:
                    count = len(self.db.find('tag', tag.tag))
                    rows.append([tag.tag, count])

                # Generate the table with the results.
                header = ['Tag', '# Entries']
                print(table(header=header, rows=rows))
            else:
                print("No tags available")

            return

        # At this point, if there are no search terms specified, return.
        if len(args) == 0:
            usage()
            return

        # The first argument is the search term (or "key").
        key = args[0]
        if key != 'all' and key != 'latest':
            try:
                # The second argument is the search value.
                value = args[1]
            except IndexError:
                print_error("You need to include a search term.")
                return
        else:
            value = None

        # Search all the files matching the given parameters.
        items = self.db.find(key, value)
        if not items:
            return

        # Populate the list of search results.
        rows = []
        count = 1
        for item in items:
            tag = ', '.join([t.tag for t in item.tag if t.tag])
            row = [count, item.name, item.mime, item.md5, tag]
            if key == 'latest':
                row.append(item.created_at)

            rows.append(row)
            count += 1

        # Update find results in current session.
        __sessions__.find = items

        # Generate a table with the results.
        header = ['#', 'Name', 'Mime', 'MD5', 'Tags']
        if key == 'latest':
            header.append('Created At')

        print(table(header=header, rows=rows))

    ##
    # TAGS
    #
    # This command is used to modify the tags of the opened file.
    def cmd_tags(self, *args):
        def usage():
            print("usage: tags [-h] [-a=tags] [-d=tag]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--add (-a)\tAdd tags to the opened file (comma separated)")
            print("\t--delete (-d)\tDelete a tag from the opened file")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'ha:d:', ['help', 'add=', 'delete='])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_add = None
        arg_delete = None

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-a', '--add'):
                arg_add = value
            elif opt in ('-d', '--delete'):
                arg_delete = value

        # This command requires a session to be opened.
        if not __sessions__.is_set():
            print_error("No session opened")
            return

        # If no arguments are specified, there's not much to do.
        # However, it could make sense to also retrieve a list of existing
        # tags from this command, and not just from the "find" command alone.
        if not arg_add and not arg_delete:
            usage()
            return

        # TODO: handle situation where addition or deletion of a tag fail.

        if arg_add:
            # Add specified tags to the database's entry belonging to
            # the opened file.
            db = Database()
            db.add_tags(__sessions__.current.file.sha256, arg_add)
            print_info("Tags added to the currently opened file")

            # We refresh the opened session to update the attributes.
            # Namely, the list of tags returned by the "info" command
            # needs to be re-generated, or it wouldn't show the new tags
            # until the existing session is closed a new one is opened.
            print_info("Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

        if arg_delete:
            # Delete the tag from the database.
            Database().delete_tag(arg_delete)
            # Refresh the session so that the attributes of the file are
            # updated.
            print_info("Refreshing session to update attributes...")
            __sessions__.new(__sessions__.current.file.path)

    ###
    # SESSION
    #
    # This command is used to list and switch across all the opened sessions.
    def cmd_sessions(self, *args):
        def usage():
            print("usage: sessions [-h] [-l] [-s=session]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--list (-l)\tList all existing sessions")
            print("\t--switch (-s)\tSwitch to the specified session")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hls:', ['help', 'list', 'switch='])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_list = False
        arg_switch = None

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-l', '--list'):
                arg_list = True
            elif opt in ('-s', '--switch'):
                arg_switch = int(value)

        if arg_list:
            if not __sessions__.sessions:
                print_info("There are no opened sessions")
                return

            rows = []
            for session in __sessions__.sessions:
                current = ''
                if session == __sessions__.current:
                    current = 'Yes'

                rows.append([
                    session.id,
                    session.file.name,
                    session.file.md5,
                    session.created_at,
                    current
                ])

            print_info("Opened Sessions:")
            print(table(header=['#', 'Name', 'MD5', 'Created At', 'Current'], rows=rows))
            return
        elif arg_switch:
            for session in __sessions__.sessions:
                if arg_switch == session.id:
                    __sessions__.switch(session)
                    return

            print_warning("The specified session ID doesn't seem to exist")
            return

        usage()

    ##
    # PROJECTS
    #
    # This command retrieves a list of all projects.
    # You can also switch to a different project.
    def cmd_projects(self, *args):
        def usage():
            print("usage: projects [-h] [-l] [-s=project]")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--list (-l)\tList all existing projects")
            print("\t--switch (-s)\tSwitch to the specified project")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hls:', ['help', 'list', 'switch='])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_list = False
        arg_switch = None

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-l', '--list'):
                arg_list = True
            elif opt in ('-s', '--switch'):
                arg_switch = value

        projects_path = os.path.join(os.getcwd(), 'projects')

        if not os.path.exists(projects_path):
            print_info("The projects directory does not exist yet")
            return

        if arg_list:
            print_info("Projects Available:")

            rows = []
            for project in os.listdir(projects_path):
                project_path = os.path.join(projects_path, project)
                if os.path.isdir(project_path):
                    current = ''
                    if __project__.name and project == __project__.name:
                        current = 'Yes'
                    rows.append([project, time.ctime(os.path.getctime(project_path)), current])

            print(table(header=['Project Name', 'Creation Time', 'Current'], rows=rows))
            return
        elif arg_switch:
            if __sessions__.is_set():
                __sessions__.close()
                print_info("Closed opened session")

            __project__.open(arg_switch)
            print_info("Switched to project {0}".format(bold(arg_switch)))

            # Need to re-initialize the Database to open the new SQLite file.
            self.db = Database()
            return

        usage()

    ##
    # EXPORT
    #
    # This command will export the current session to file or zip.
    def cmd_export(self, *args):
        def usage():
            print("usage: export [-h] [-z] <path or archive name>")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--zip (-z)\tExport session in a zip archive")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hz', ['help', 'zip'])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        arg_zip = False

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-z', '--zip'):
                arg_zip = True
                
        # This command requires a session to be opened.
        if not __sessions__.is_set():
            print_error("No session opened")
            return

        # Check for valid export path.
        if len(args) ==0:
            usage()
            return

        # TODO: having for one a folder and for the other a full
        # target path can be confusing. We should perhaps standardize this.

        # Abort if the specified path already exists.
        if os.path.isfile(argv[0]):
            print_error("File at path \"{0}\" already exists, abort".format(argv[0]))
            return

        # If the argument chosed so, archive the file when exporting it.
        # TODO: perhaps add an option to use a password for the archive
        # and default it to "infected".
        if arg_zip:
            try:
                with ZipFile(argv[0], 'w') as export_zip:
                    export_zip.write(__sessions__.current.file.path, arcname=__sessions__.current.file.name)
            except IOError as e:
                print_error("Unable to export file: {0}".format(e))
            else:
                print_info("File archived and exported to {0}".format(argv[0]))
        # Otherwise just dump it to the given directory.
        else:
            # XXX: Export file with the original file name.
            store_path = os.path.join(argv[0], __sessions__.current.file.name)

            try:
                shutil.copyfile(__sessions__.current.file.path, store_path)
            except IOError as e:
                print_error("Unable to export file: {0}".format(e))
            else:
                print_info("File exported to {0}".format(store_path))