Python idaapi.get_root_filename() Examples

def log_all():
    data = {}
    idc.auto_wait()
    analyser = Analyser()
    if not analyser.valid:
        idc.qexit(-1)
    analyser.get_boot_services()
    module = idaapi.get_root_filename()
    boot_services = get_boot_services(analyser)
    protocols = get_protocols(analyser)
    data['module_name'] = module
    data['boot_services'] = boot_services
    data['protocols'] = protocols
    logs_dir = os.path.join(tempfile.gettempdir(), 'uefi-retool-all-info')
    if not os.path.isdir(logs_dir):
        os.mkdir(logs_dir)
    log_fname = os.path.join(
        logs_dir, '{}.json'.format(
            binascii.hexlify(ida_nalt.retrieve_input_file_md5()).decode()))
    with open(log_fname, 'w') as f:
        json.dump(data, f, indent=4)
    idc.qexit(0) 

def on_open():
    global auto_wait
    global fhash
    print('revsync: file opened:', idaapi.get_root_filename())
    netnode.create(NETNODE_NAME)
    try: fhash = netnode.getblob(0, 'I').decode('ascii')
    except: fhash = None
    if not fhash:
        fhash = read_fhash()
        try: ret = netnode.setblob(fhash.encode('ascii'), 0, 'I')
        except: print('saving fhash failed, this will probably break revsync')

    if auto_is_ok():
        on_load()
        auto_wait = False
    else:
        auto_wait = True
        print('revsync: waiting for auto analysis')
        if not hasattr(IDP_Hooks, 'auto_empty_finally'):
            idaapi.register_timer(1000, wait_for_analysis) 

def get_all_pooltags():
	""" Returns a string with a 'pooltags.txt' formatted string of 'pool tag' - 'driver' - 'functions which use it'.
	"""
	
	tags = find_pool_tags()
	out = ''
	file_name = idaapi.get_root_filename()
	for tag in tags.keys():
		desc = 'Called by: '
		desc += ', '.join(tags[tag])
		out += '{} - {} - {}\n'.format(tag, file_name, desc)
	return out 

def do_export():
    db = {}
    module = idaapi.get_root_filename().lower()
    base = idaapi.get_imagebase()

    file = ida_kernwin.ask_file(1, "x64dbg database|{}".format(get_file_mask()),
                                "Export database")
    if not file:
        return
    print("Exporting database {}".format(file))

    db["labels"] = [{
        "text": name,
        "manual": False,
        "module": module,
        "address": "{:#x}".format(ea - base)
    } for (ea, name) in idautils.Names()]
    print("{:d} label(s) exported".format(len(db["labels"])))

    db["comments"] = [{
        "text": comment.replace("{", "{{").replace("}", "}}"),
        "manual": False,
        "module": module,
        "address": "{:#x}".format((ea - base))
    } for (ea, comment) in Comments()]
    print("{:d} comment(s) exported".format(len(db["comments"])))

    db["breakpoints"] = [{
        "address": "{:#x}".format(ea - base),
        "enabled": True,
        "type": bptype,
        "titantype": "{:#x}".format(titantype),
        "oldbytes": "{:#x}".format(oldbytes),
        "module": module,
    } for (ea, bptype, titantype, oldbytes) in Breakpoints()]
    print("{:d} breakpoint(s) exported".format(len(db["breakpoints"])))

    with open(file, "w") as outfile:
        json.dump(db, outfile, indent=1)
    print("Done!") 

def extract(self):
        """Extract the control flow graph from the binary."""
        # Allocate a new graph
        self.graph = graph_alloc(0)
        
        # Initialize binary info
        self.info = get_inf_structure()
        
        # Initialize Capstone
        if self.info.is_64bit():
            mode = capstone.CS_MODE_64
        else:
            mode = capstone.CS_MODE_32
        self.capstone = capstone.Cs(capstone.CS_ARCH_X86, mode)
        
        # Get the Entry Point
        entry = None
        try:
            start_ea = self.info.start_ea
            if start_ea != 0xffffffff:
                entry = start_ea
        except:
            try:
                entry = BeginEA()
            except:
                pass
                
        if entry is None:
            print("WARNING: Could not determine entrypoint")
        else:
            self.dis(ea=entry, is_child1=None, ifrom=None)

        # Scan all the functions
        for ea in Functions():
            self.dis(ea=ea, is_child1=None, ifrom=None)

        update_children_fathers_number(self.graph)

        # Information
        print("%s graph has %d nodes" % (get_root_filename(),
                                         self.graph.nodes.size)) 

def get_root_filename(self):
        return idaapi.get_root_filename() 

def log_pp_guids():
    idc.auto_wait()
    analyser = Analyser()
    if not analyser.valid:
        idc.qexit(-1)
    analyser.get_boot_services()
    analyser.get_protocols()
    analyser.get_prot_names()
    data = {}
    data['module_name'] = idaapi.get_root_filename()
    data['protocols'] = []
    for protocol_record in analyser.Protocols['all']:
        if (protocol_record['protocol_name'] == 'ProprietaryProtocol'):
            guid = get_guid_str(protocol_record['guid'])
            service = protocol_record['service']
            address = '{addr:#x}'.format(addr=protocol_record['address'])
            data['protocols'].append({
                'guid': guid,
                'service': service,
                'address': address
            })
    logs_dir = os.path.join(tempfile.gettempdir(), 'uefi-retool-pp-guids')
    if not os.path.isdir(logs_dir):
        os.mkdir(logs_dir)
    log_fname = os.path.join(
        logs_dir, '{}.json'.format(
            binascii.hexlify(ida_nalt.retrieve_input_file_md5()).decode()))
    with open(log_fname, 'w') as f:
        json.dump(data, f, indent=4)
    idc.qexit(0) 

def read_fhash():
    filename = idaapi.get_root_filename()
    if filename is None:
        return None
    with open(filename, 'rb') as f:
        return hashlib.sha256(f.read()).hexdigest().upper() 

def setup():
    if idaapi.get_root_filename():
        on_open()
    else:
        idaapi.notify_when(idaapi.NW_OPENIDB | idaapi.NW_CLOSEIDB | idaapi.NW_TERMIDA, eventhook)
    return -1 

def upload(self,ctx):
		start = time.time()
		func_count = 0
		bb_count = 0
		call_count = 0
		target = idaapi.get_root_filename()
		hash = idc.GetInputMD5()
		tx = self.neo.cypher.begin()
		insert_binary = "MERGE (n:Binary {name:{N},hash:{H}}) RETURN n"
		insert_func = "MERGE (n:Function {name:{N},start:{S},flags:{F}}) RETURN n"
		insert_bb = "MERGE (n:BasicBlock {start:{S}, end:{E}}) RETURN n"
		create_relationship = "MATCH (u:Function {name:{N}}), (r:Function {start:{S}}) CREATE (u)-[:CALLS]->(r)"
		create_contains = "MATCH (u:BasicBlock {start:{S}}), (f:Function {name:{N}}) CREATE (f)-[:CONTAINS]->(u)"
		create_inside = "MATCH (u:Function {start:{S}}), (b:Binary {hash:{H}}) CREATE (f)-[:INSIDE]->(b)"
		self.neo.cypher.execute(insert_binary, {"N":target, "H":hash})
		self.neo.cypher.execute("CREATE INDEX ON :Function(start)")
		#self.neo.cypher.execute("CREATE INDEX ON :Function(name)")
		self.neo.cypher.execute("CREATE INDEX ON :BasicBlock(start)")
		for f in Functions():
			tx.append(create_inside, {"S":f, "H":hash})
			callee_name = GetFunctionName(f)
			flags = get_flags(f)
			type = GetType(f)
			if type:
				return_type = type.split()[0]
				print type
				end_return = type.find(' ')
				start_args = type.find('(')
				print type[end_return +1:start_args]
				print type[start_args+1:].split(',')
			else:
				print GuessType(f)
			tx.append(insert_func, {"N": callee_name, "S":f, "F":flags})
			func_count += 1
			fc = idaapi.FlowChart(idaapi.get_func(f))
			for block in fc:
				tx.append(insert_bb, {"S":block.startEA,"E":block.endEA})
				tx.append(create_contains,{"S":block.startEA,"N":f})
				bb_count += 1
		tx.process()
		tx.commit()
		tx = self.neo.cypher.begin()
		for f in Functions():
			for xref in CodeRefsTo(f,0):
				caller_name = GetFunctionName(xref)
				if caller_name != '':
					tx.append(create_relationship,{"N":caller_name,"S":f})
					call_count += 1
		tx.process()
		tx.commit()
		print "Upload ran in: " + str(time.time() - start)
		print "Uploaded " + str(func_count) + " functions, " + str(call_count) +" function calls and " + str(bb_count) + " basic blocks."