Python __sessions__.close函数代码示例

    def run(self, *args):
        try:
            self.parser.parse_args(args)
        except SystemExit:
            return

        __sessions__.close()

def module_cmdline(cmd_line, file_hash):
    html = ""
    cmd = Commands()
    split_commands = cmd_line.split(';')
    for split_command in split_commands:
        split_command = split_command.strip()
        if not split_command:
            continue
        root, args = parse(split_command)
        try:
            if root in cmd.commands:
                cmd.commands[root]['obj'](*args)
                html += print_output(cmd.output)
                del (cmd.output[:])
            elif root in __modules__:
                # if prev commands did not open a session open one on the current file
                if file_hash:
                    path = get_sample_path(file_hash)
                    __sessions__.new(path)
                module = __modules__[root]['obj']()
                module.set_commandline(args)
                module.run()

                html += print_output(module.output)
                if cfg.modules.store_output and __sessions__.is_set():
                    Database().add_analysis(file_hash, split_command, module.output)
                del (module.output[:])
            else:
                html += '<p class="text-danger">{0} is not a valid command</p>'.format(cmd_line)
        except Exception as e:
            html += '<p class="text-danger">We were unable to complete the command {0}</p>'.format(cmd_line)
    __sessions__.close()
    return html

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

    def run(self, *args):
        try:
            args = self.parser.parse_args(args)
        except SystemExit:
            return

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

        db = Database()

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

            samples = db.find('all')
            for sample in samples:
                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:
                    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 = db.find('sha256', __sessions__.current.file.sha256)
                if rows:
                    malware_id = rows[0].id
                    if 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.")

    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

        cfg = Config()
        if cfg.paths.storage_path:
            base_path = cfg.paths.storage_path
        else:
            base_path = os.path.join(expanduser("~"), '.viper')

        projects_path = os.path.join(base_path, '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())

def file_info(file_hash, project=False):
    contents = {}
    if project in project_list():
        __project__.open(project)
        contents['project'] = project
    else:
        __project__.open('../')
        contents['project'] = 'Main'
    # Open the Database
    db = Database()
    # Open a session
    try:
        path = get_sample_path(file_hash)
        __sessions__.new(path)
    except:
        return template('error.tpl', error="{0} Does not match any hash in the Database".format(file_hash))
    
    # Get the file info
    contents['file_info'] = [
                __sessions__.current.file.name,
                __sessions__.current.file.tags,
                __sessions__.current.file.path,
                __sessions__.current.file.size,
                __sessions__.current.file.type,
                __sessions__.current.file.mime,
                __sessions__.current.file.md5,
                __sessions__.current.file.sha1,
                __sessions__.current.file.sha256,
                __sessions__.current.file.sha512,
                __sessions__.current.file.ssdeep,
                __sessions__.current.file.crc32                
                ]
                
    # Get Any Notes
    note_list = []
    malware = db.find(key='sha256', value=file_hash)
    if malware:
        notes = malware[0].note
        if notes:
            rows = []
            for note in notes:
                note_list.append([note.title, note.body, note.id])
    contents['notes'] = note_list
    
    # Close the session
    __sessions__.close()
    # Return the page
    return template('file.tpl', **contents)

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

def add_file(file_path, tags, parent):
    obj = File(file_path)
    new_path = store_sample(obj)
    print new_path
    success = True
    if new_path:
        # Add file to the database.
        db = Database()
        success = db.add(obj=obj, 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

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

def run_module():
    # Optional Project
    project = request.forms.get('project')
    # Optional sha256   
    sha256 = request.forms.get('sha256')
    # Not Optional Command Line Args
    cmd_line = request.forms.get('cmdline')
    if project:
        __project__.open(project)
    else:
        __project__.open('../')
    if sha256:
        file_path = get_sample_path(sha256)
        if file_path:
            __sessions__.new(file_path)
    if not cmd_line:
        response.code = 404
        return jsonize({'message':'Invalid command line'})
    results = module_cmdline(cmd_line, sha256)  
    __sessions__.close()
    return jsonize(results)

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

def module_text(file_hash, cmd_string):
    # A lot of commands rely on an open session
    # open a session on the file hash
    path = get_sample_path(file_hash)
    __sessions__.new(path)
    # Run the Module with args
    if __sessions__.is_set():
        root, args = parse(cmd_string)
        if root in __modules__:
            module = __modules__[root]['obj']()
            module.set_args(args)
            module.run()
            html = print_output(module.output)
            del(module.output[:])
        else:
            html = '<p class="text-danger">{0} is not a valid command</p>'.format(cmd_string)
    else:
        '<p class="text-danger">! There is no open session</p>'
    # close the session
    __sessions__.close()
    return html

def module_cmdline(cmd_line, sha256):
    json_data = ''
    cmd = Commands()
    split_commands = cmd_line.split(';')
    for split_command in split_commands:
        split_command = split_command.strip()
        if not split_command:
            continue
        root = ''
        args = []
        # Split words by white space.
        words = split_command.split()
        # First word is the root command.
        root = words[0]
        # If there are more words, populate the arguments list.
        if len(words) > 1:
            args = words[1:]
        try:
            if root in cmd.commands:
                cmd.commands[root]['obj'](*args)
                if cmd.output:
                    json_data += str(cmd.output)
                del(cmd.output[:])
            elif root in __modules__:
                # if prev commands did not open a session open one on the current file
                if sha256:
                    path = get_sample_path(sha256)
                    __sessions__.new(path)
                module = __modules__[root]['obj']()
                module.set_commandline(args)
                module.run()

                json_data += str(module.output)
                del(module.output[:])
            else:
                json_data += "{'message':'{0} is not a valid command'.format(cmd_line)}"
        except:
            json_data += "{'message':'Unable to complete the command: {0}'.format(cmd_line)}"
    __sessions__.close()
    return json_data 

def module_text(file_hash, cmd_string):
    # redirect stdout to a memory file
    sys.stdout = StringIO()

    # A lot of commands rely on an open session
    # open a session on the file hash
    path = get_sample_path(file_hash)
    __sessions__.new(path)
    
    # Run the Module with args, the ouptut will end up in a StringIO Object
    if __sessions__.is_set():
        root, args = parse(cmd_string)
        if root in __modules__:
            module = __modules__[root]['obj']()
            module.set_args(args)
            module.run()
    
    # Parse the raw data to something we can use
    clean_text = parse_text(sys.stdout.getvalue())
    # close the session
    __sessions__.close()
    return clean_text

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