Python Database.add方法代码示例

# 需要导入模块: from viper.core.database import Database [as 别名]
# 或者: from viper.core.database.Database import add [as 别名]
 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))

# 需要导入模块: from viper.core.database import Database [as 别名]
# 或者: from viper.core.database.Database import add [as 别名]
 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

# 需要导入模块: from viper.core.database import Database [as 别名]
# 或者: from viper.core.database.Database import add [as 别名]
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

# 需要导入模块: from viper.core.database import Database [as 别名]
# 或者: from viper.core.database.Database import add [as 别名]
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")

# 需要导入模块: from viper.core.database import Database [as 别名]
# 或者: from viper.core.database.Database import add [as 别名]
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

# 需要导入模块: from viper.core.database import Database [as 别名]
# 或者: from viper.core.database.Database import add [as 别名]
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")

# 需要导入模块: from viper.core.database import Database [as 别名]
# 或者: from viper.core.database.Database import add [as 别名]
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
#.........这里部分代码省略.........

# 需要导入模块: from viper.core.database import Database [as 别名]
# 或者: from viper.core.database.Database import add [as 别名]
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"
        )
#.........这里部分代码省略.........

# 需要导入模块: from viper.core.database import Database [as 别名]
# 或者: from viper.core.database.Database import add [as 别名]
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")
#.........这里部分代码省略.........

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

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

        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))

# 需要导入模块: from viper.core.database import Database [as 别名]
# 或者: from viper.core.database.Database import add [as 别名]
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
#.........这里部分代码省略.........

# 需要导入模块: from viper.core.database import Database [as 别名]
# 或者: from viper.core.database.Database import add [as 别名]
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))

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

#.........这里部分代码省略.........
        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)

# 需要导入模块: from viper.core.database import Database [as 别名]
# 或者: from viper.core.database.Database import add [as 别名]
    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