def embedd_configuration(db_name, model, batch_size, max_instruction, embeddings_table): msg = "Embedding options: \n" msg += " - Database Name: {} \n".format(db_name) msg += " - Model: {} \n".format(model) msg += " - Batch Size: {} \n".format(batch_size) msg += " - Max Instruction per function: {} \n".format(max_instruction) msg += " - Table for saving embeddings: {}.".format(embeddings_table) return msg if __name__ == '__main__': print_safe() parser = argparse.ArgumentParser(description=debug_msg) parser.add_argument("-db", "--db", help="Name of the database to create", required=True) parser.add_argument("-b", "--build", help="Build db disassebling executables", action="store_true") parser.add_argument("-s", "--split", help="Perform data splitting for training",
for function in functions: if functions[function]['address'] == address: instructions_list = functions[function]['filtered_instructions'] break if instructions_list is None: print("Function not found") return None converted_instructions = self.converter.convert_to_ids(instructions_list) instructions, length = self.normalizer.normalize_functions([converted_instructions]) embedding = self.embedder.embedd(instructions, length) return embedding if __name__ == '__main__': utils.print_safe() parser = ArgumentParser(description="Safe Embedder") parser.add_argument("-m", "--model", help="Safe trained model to generate function embeddings") parser.add_argument("-i", "--input", help="Input executable that contains the function to embedd") parser.add_argument("-a", "--address", help="Hexadecimal address of the function to embedd") args = parser.parse_args() address = int(args.address, 16) safe = SAFE(args.model) embedding = safe.embedd_function(args.input, address) print(embedding[0])