process_firmware.py 文件源码

def register_files(self):
        print("Start registering files")
        for root, dirs, files in os.walk(self.extracted_path):
            for file in files:
                full_path = os.path.join(root, file)
                if not os.path.isfile(full_path):
                    continue
                path = full_path.replace(self.extracted_path, "")
                content = ""
                hash = ""
                with open(full_path, "rb") as fd:
                    content = fd.read()
                    hash_content = "%s:%s" % (file, content)
                    hash = hashlib.md5(hash_content.encode('utf-8')).hexdigest()
                try:
                    file_obj = FileModel.objects.get(hash=hash)
                    file_obj.firmware.add(self.firmware)
                    file_obj.save()
                except FileModel.DoesNotExist:
                    try:
                        file_obj = FileModel()
                        file_obj.filepath = os.path.join(root, file)
                        file_obj.hash = hash
                        file_obj.filesize = len(content)
                        file_obj.filename = path
                        file_obj.save()
                        file_obj.firmware.add(self.firmware)
                        file_obj.file_type = magic.from_file(os.path.join(root,
                                                                          file))
                        file_obj.save()
                        self.find_loots(file_obj)
                        # Performance tweak
                        file_obj.nb_loots = file_obj.loots.all().count()
                    except:
                        file_obj.file_type = "unknown"

        print("Files registered")