Python database.Database类代码示例

    def run(self):
        if not __session__.is_set():
            print_error("No session opened")
            return

        if not HAVE_PYDEEP:
            print_error("Missing dependency, install pydeep (`pip install pydeep`)")
            return

        if not __session__.file.ssdeep:
            print_error("No ssdeep hash available for opened file")
            return

        db = Database()
        samples = db.find(key='all')

        for sample in samples:
            if sample.sha256 == __session__.file.sha256:
                continue

            if not sample.ssdeep:
                continue

            score = pydeep.compare(__session__.file.ssdeep, sample.ssdeep)
            if score > 40:
                print("Match {0}%: {1}".format(score, sample.sha256))

    def size_cluster(self):
        db = Database()
        samples = db.find(key='all')

        cluster = {}
        for sample in samples:
            sample_path = get_sample_path(sample.sha256)
            if not os.path.exists(sample_path):
                continue

            try:
                cur_size = os.path.getsize(sample_path)
            except Exception as e:
                self.log('error', "Error {0} for sample {1}".format(e, sample.sha256))
                continue

            if cur_size not in cluster:
                cluster[cur_size] = []

            cluster[cur_size].append([sample.md5, sample.name])

        for cluster_name, cluster_members in cluster.items():
            # Skipping clusters with only one entry.
            if len(cluster_members) == 1:
                continue

            self.log('info', "Cluster size {0} with {1} elements".format(bold(cluster_name), len(cluster_members)))
            self.log('table', dict(header=['MD5', 'Name'], rows=cluster_members))

def delete_file(file_hash):
    success = False
    key = ''
    if len(file_hash) == 32:
        key = 'md5'
    elif len(file_hash) == 64:
        key = 'sha256'
    else:
        return HTTPError(400, 'Invalid hash format (use md5 or sha256)')

    db = Database()
    rows = db.find(key=key, value=file_hash)

    if not rows:
        raise HTTPError(404, 'File not found in the database')
	
    if rows:
		malware_id = rows[0].id
		path = get_sample_path(rows[0].sha256)
		if db.delete(malware_id):
			success = True
		else:
			raise HTTPError(404, 'File not found in repository')

    path = get_sample_path(rows[0].sha256)
    if not path:
        raise HTTPError(404, 'File not found in file system')
    else:	
        success=os.remove(path)
    
    if success:
        return jsonize({'message' : 'deleted'})
    else:
        return HTTPError(500, 'Unable to delete file')

    def run(self):
        super(Strings, self).run()
        
        if self.args is None:
            return

        if not (self.args.all or self.args.files or self.args.hosts or self.args.network or self.args.interesting):
            self.log('error', 'At least one of the parameters is required')
            self.usage()
            return

        if self.args.scan:
            db = Database()
            samples = db.find(key='all')
            for sample in samples:
                sample_path = get_sample_path(sample.sha256)
                strings = self.get_strings(File(sample_path))
                self.process_strings(strings, sample.name)
        else:
            if not __sessions__.is_set():
                self.log('error', "No open session")
                return
            if os.path.exists(__sessions__.current.file.path):
                strings = self.get_strings(__sessions__.current.file)
                self.process_strings(strings)

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

def get_file(file_hash):
    key = ''
    if len(file_hash) == 32:
        key = 'md5'
    elif len(file_hash) == 64:
        key = 'sha256'
    else:
        return HTTPError(400, 'Invalid hash format (use md5 or sha256)')

    db = Database()
    rows = db.find(key=key, value=file_hash)

    if not rows:
        raise HTTPError(404, 'File not found in the database')

    path = get_sample_path(rows[0].sha256)
    if not path:
        raise HTTPError(404, 'File not found in the repository')

    response.content_length = os.path.getsize(path)
    response.content_type = 'application/octet-stream; charset=UTF-8'
    data = ''
    for chunk in File(path).get_chunks():
        data += chunk

    return data

def logo():
    print("""         _
        (_)
   _   _ _ ____  _____  ____
  | | | | |  _ \| ___ |/ ___)
   \ V /| | |_| | ____| |
    \_/ |_|  __/|_____)_| v{}
          |_|
    """.format(__version__))

    db = Database()
    count = db.get_sample_count()

    try:
        db.find('all')
    except Exception:
        print_error("You need to update your Viper database. Run 'python update.py -d'")
        sys.exit()

    if __project__.name:
        name = __project__.name
    else:
        name = 'default'

    print(magenta("You have " + bold(count)) +
          magenta(" files in your " + bold(name)) +
          magenta(" repository"))

def logo():
    print("""         _
        (_)
   _   _ _ ____  _____  ____
  | | | | |  _ \| ___ |/ ___)
   \ V /| | |_| | ____| |
    \_/ |_|  __/|_____)_| v1.3-dev
          |_|
    """)

    db = Database()
    count = db.get_sample_count()

    # Handle the New database format
    try:
        db.find('all', None)
    except:
        print_error("You need to update your viper database. Run 'python update.py -d'")
        sys.exit()


    if __project__.name:
        name = __project__.name
    else:
        name = 'default'

    print(magenta("You have " + bold(count)) +
          magenta(" files in your " + bold(name) +
          magenta(" repository".format(bold(name)))))
    if cfg.autorun.enabled and len(cfg.autorun.commands) == 0:
        print_warning("You have enabled autorun but not set any commands in viper.conf.")

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

        if not __sessions__.is_set():
            self.log('error', "No open session. This command expects a file to be open.")
            return

        if not __project__.name:
            src_project = "default"
        else:
            src_project = __project__.name

        db = Database()

        db.copied_id_sha256 = []
        res = db.copy(__sessions__.current.file.id,
                      src_project=src_project, dst_project=args.project,
                      copy_analysis=True, copy_notes=True, copy_tags=True, copy_children=args.children)

        if args.delete:
            __sessions__.close()
            for item_id, item_sha256 in db.copied_id_sha256:
                db.delete_file(item_id)
                os.remove(get_sample_path(item_sha256))
                self.log('info', "Deleted: {}".format(item_sha256))

        if res:
            self.log('success', "Successfully copied sample(s)")
            return True
        else:
            self.log('error', "Something went wrong")
            return False

 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

    def peid(self):

        def get_signatures():
            with file(os.path.join(VIPER_ROOT, 'data/peid/UserDB.TXT'), 'rt') as f:
                sig_data = f.read()

            signatures = peutils.SignatureDatabase(data=sig_data)

            return signatures

        def get_matches(pe, signatures):
            matches = signatures.match_all(pe, ep_only=True)
            return matches

        if not self.__check_session():
            return

        signatures = get_signatures()
        peid_matches = get_matches(self.pe, signatures)

        if peid_matches:
            self.log('info', "PEiD Signatures:")
            for sig in peid_matches:
                if type(sig) is list:
                    self.log('item', sig[0])
                else:
                    self.log('item', sig)
        else:
            self.log('info', "No PEiD signatures matched.")

        if self.args.scan and peid_matches:
            self.log('info', "Scanning the repository for matching samples...")

            db = Database()
            samples = db.find(key='all')

            matches = []
            for sample in samples:
                if sample.sha256 == __sessions__.current.file.sha256:
                    continue

                sample_path = get_sample_path(sample.sha256)
                if not os.path.exists(sample_path):
                    continue

                try:
                    cur_pe = pefile.PE(sample_path)
                    cur_peid_matches = get_matches(cur_pe, signatures)
                except:
                    continue

                if peid_matches == cur_peid_matches:
                    matches.append([sample.name, sample.sha256])

            self.log('info', "{0} relevant matches found".format(bold(len(matches))))

            if len(matches) > 0:
                self.log('table', dict(header=['Name', 'SHA256'], rows=matches))

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

    def pehash(self):
        if not HAVE_PEHASH:
            self.log('error', "PEhash is missing. Please copy PEhash to the modules directory of Viper")
            return

        current_pehash = None
        if __sessions__.is_set():
            current_pehash = calculate_pehash(__sessions__.current.file.path)
            self.log('info', "PEhash: {0}".format(bold(current_pehash)))

        if self.args.all or self.args.cluster or self.args.scan:
            db = Database()
            samples = db.find(key='all')

            rows = []
            for sample in samples:
                sample_path = get_sample_path(sample.sha256)
                pe_hash = calculate_pehash(sample_path)
                if pe_hash:
                    rows.append((sample.name, sample.md5, pe_hash))

        if self.args.all:
            self.log('info', "PEhash for all files:")
            header = ['Name', 'MD5', 'PEhash']
            self.log('table', dict(header=header, rows=rows))

        elif self.args.cluster:
            self.log('info', "Clustering files by PEhash...")

            cluster = {}
            for sample_name, sample_md5, pe_hash in rows:
                cluster.setdefault(pe_hash, []).append([sample_name, sample_md5])

            for item in cluster.items():
                if len(item[1]) > 1:
                    self.log('info', "PEhash cluster {0}:".format(bold(item[0])))
                    self.log('table', dict(header=['Name', 'MD5'], rows=item[1]))

        elif self.args.scan:
            if __sessions__.is_set() and current_pehash:
                self.log('info', "Finding matching samples...")

                matches = []
                for row in rows:
                    if row[1] == __sessions__.current.file.md5:
                        continue

                    if row[2] == current_pehash:
                        matches.append([row[0], row[1]])

                if matches:
                    self.log('table', dict(header=['Name', 'MD5'], rows=matches))
                else:
                    self.log('info', "No matches found")

    def compiletime(self):

        def get_compiletime(pe):
            return datetime.datetime.fromtimestamp(pe.FILE_HEADER.TimeDateStamp)

        if not self.__check_session():
            return

        compile_time = get_compiletime(self.pe)
        self.log('info', "Compile Time: {0}".format(bold(compile_time)))

        if self.args.scan:
            self.log('info', "Scanning the repository for matching samples...")

            db = Database()
            samples = db.find(key='all')

            matches = []
            for sample in samples:
                if sample.sha256 == __sessions__.current.file.sha256:
                    continue

                sample_path = get_sample_path(sample.sha256)
                if not os.path.exists(sample_path):
                    continue

                try:
                    cur_pe = pefile.PE(sample_path)
                    cur_compile_time = get_compiletime(cur_pe)
                except:
                    continue

                if compile_time == cur_compile_time:
                    matches.append([sample.name, sample.md5, cur_compile_time])
                else:
                    if self.args.window:
                        if cur_compile_time > compile_time:
                            delta = (cur_compile_time - compile_time)
                        elif cur_compile_time < compile_time:
                            delta = (compile_time - cur_compile_time)

                        delta_minutes = int(delta.total_seconds()) / 60
                        if delta_minutes <= self.args.window:
                            matches.append([sample.name, sample.md5, cur_compile_time])

            self.log('info', "{0} relevant matches found".format(bold(len(matches))))

            if len(matches) > 0:
                self.log('table', dict(header=['Name', 'MD5', 'Compile Time'], rows=matches))

def tags():
    # Set DB
    db = Database()
    
    # Search or Delete
    if request.method == 'GET':
        action = request.query.action
        value = request.query.value.strip()
        
        if value:
            if action == 'delete':
                # Delete individual tags is not in viper yet
                pass
            elif action == 'search':
                # This will search all projects
                # Get project list
                projects = project_list()
                # Add Main db to list.
                projects.append('../')
                # Search All projects
                p_list = []
                results = {}
                for project in projects:
                    __project__.open(project)
                    # Init DB
                    db = Database()
                    #get results
                    proj_results = []
                    rows = db.find(key='tag', value=value)
                    for row in rows:
                        if project == '../':
                            project = 'Main'
                        proj_results.append([row.name, row.sha256])
                    results[project] = proj_results
                    p_list.append(project)
                # Return the search template
                return template('search.tpl', projects=p_list, results=results)
            else:
                return template('error.tpl', error="'{0}' Is not a valid tag action".format(action))
                             
    # Add New Tags
    if request.method == 'POST':
        file_hash = request.forms.get('sha256')
        project = request.forms.get('project')
        if file_hash and project:
            tags = request.forms.get('tags')
            db.add_tags(file_hash, tags)
            redirect('/file/{0}/{1}'.format(project, file_hash))