import subprocess
import os
import re
import sys

class Read_ELF(object):
	'''
	This class is used to read ELF file
	'''
	def __init__(self, elf_path):
		# write debuginfo into file
		subprocess.run("readelf -w %s > data"%elf_path, shell=True)
		with open('data', 'rb') as f:
			self.content = f.read().decode()
		
	# get the function scope from debuginfo
	def _get_func_scope(self, func):
		self.func_name = func
		# The scope between <1>.....<1> is the function scope.
		pattern = r'(: %s\n.*?<1>)'%self.func_name
		re_text = re.findall(pattern, self.content, re.I|re.S)
		
		for text in re_text:
			# function with parameter
			if text.find('<2>') >= 0:
				#print(text)
				return text
			# function without parameter
			elif text.find('DW_AT_type') >=0 :
				return text
			
		print('%s is not in the debuginfo file!'%func)
		return -1
		
		
	# basic type mapping
	def _base_type_mapping(self, type):
		type_list = type.split(' ')
		if len(type_list) >= 3 and type_list[-1] == 'int':
			if type_list[-2] == 'signed' or type_list[-2] == 'unsigned':
				value = type_list[-2]
				for t in type_list[:-2]:
					value = value + ' ' + t
				return value
			else:
				return type
		elif len(type_list) == 2 and type_list[-1] == 'int':
			if type_list[0] == 'short' or type_list[0] == 'long':
				return type_list[0]
			else:
				return type
		else:
			return type
			
	def get_types_and_names_for_subroutine(self, subroutine):
		subroutine_param_range = re.findall(r'DW_TAG_formal_parameter.*?<2>', subroutine, re.I|re.S)
		if len(subroutine_param_range)==0:
			subroutine_list = '(void)'
			return subroutine_list
		else:
			subroutine_list = '('
			for subroutine_param in subroutine_param_range:
				#print(subroutine_param)
				if subroutine_param.find('<3>')>=0:
					continue
				#dw_at_name = re.findall(r'DW_AT_name.*: (.+)\n', subroutine_param, re.I)[0]
				dw_at_name = re.findall(r'DW_AT_name.*: (.+)\n', subroutine_param, re.I)
				dw_at_type = '><' + re.findall(r'DW_AT_type.*: <0x(.+)\n', subroutine_param, re.I)[0]
				#print(subroutine_param)
				if len(dw_at_name) == 0:
					subroutine_list = subroutine_list + self.get_value_at_location(dw_at_type)[0] + ', '
				else:
					name = dw_at_name[0]
					type = self.get_value_at_location(dw_at_type)[0]
					if type[-1] == '*':
						subroutine_list = subroutine_list + type + name + ', '
					else:
						subroutine_list = subroutine_list + type + ' ' + name + ', '
			return subroutine_list[:-2] + ')'
			
	def _get_value_for_subroutine(self, subroutine_list):
		value = ' ('
		if len(subroutine_list) == 0:
			value = value + 'void)'
		elif len(subroutine_list) == 1 and subroutine_list[0]['dw_at_name'] == 'void':
			value =value + 'void)'
		else:
			for func in subroutine_list:
				param_value, _ = self.get_value_at_location(func['dw_at_type'])
				#print(param_value, ' ', func['dw_at_name'])
				if param_value[-1] == '*':
					value = value + param_value + func['dw_at_name'] + ', '
				else:
					value = value + param_value + ' ' + func['dw_at_name'] + ', '
			value = value[:-2] + ')'
		return value
	
	# get type at location
	def get_value_at_location(self, location, key='name'):
		func_param_range = re.findall(r'%s.*?<1>'%location, self.content, re.I|re.S)[0]
		type = re.findall(r'(\(.*\))\n', func_param_range, re.I)
		if len(type) == 0:
			return -1, -1
		else:
			type = type[0]
			
			
		if type == '(DW_TAG_pointer_type)':
			value = '*'
			next_location = re.findall(r'DW_AT_type.*: <0x(.+)\n', func_param_range, re.I)
			if len(next_location) == 0:
				return 'void *', -1
			else:
				next_location = '><' + next_location[0]
			next_value, subroutine_parameter = self.get_value_at_location(next_location)
			if next_value[-1]=='*':
				#print(next_value+value)
				return next_value + value, subroutine_parameter
			elif next_value[-1]==')':
				return next_value[:-1] + value + ')', subroutine_parameter
			else:
				return next_value + ' ' + value, subroutine_parameter
		elif type == '(DW_TAG_const_type)':
			#print(func_param_range)
			next_location = re.findall(r'DW_AT_type.*: <0x(.+)\n', func_param_range, re.I)
			if len(next_location) == 0:
				return 'const void', -1
			else:
				next_location = '><' + next_location[0]
				next_value, subroutine_parameter = self.get_value_at_location(next_location)
				return 'const ' + next_value, subroutine_parameter
		elif type == '(DW_TAG_base_type)':
			value = re.findall(r'DW_AT_name.*: (.+)\n', func_param_range, re.I)[0]
			value = self._base_type_mapping(value)
			return value, -1
		elif type == '(DW_TAG_subroutine_type)':
			if func_param_range.find('<2>') >= 0:
				next_location = re.findall(r'DW_AT_type.*: <0x(.+)\n', re.findall(r'.*?<2>', func_param_range, re.I|re.S)[0], re.I)
			else:
				next_location = re.findall(r'DW_AT_type.*: <0x(.+)\n', func_param_range, re.I)
			subroutine_parameter = self.get_types_and_names_for_subroutine(func_param_range)
			if len(next_location) == 0:
				return 'void ()', subroutine_parameter
			else:
				next_location = '><' + next_location[0]
				next_value, _ = self.get_value_at_location(next_location)
				#print('** ', next_value, subroutine_parameter)
				return next_value + ' ()', subroutine_parameter
		else:
			value = re.findall(r'DW_AT_name.*: (.+)\n', func_param_range, re.I)[0]
			return value, -1
	
	
	def get_types_and_names_for_func(self, func):
		self.func_range = self._get_func_scope(func)
		if self.func_range == -1:
			return -1
		#print(self.func_range)
		self.func_list = []
		self.func_param_range = re.findall(r'DW_TAG_formal_parameter.*?<2>', self.func_range, re.I|re.S)
		if len(self.func_param_range)==0:
			self.no_parameter = True
			self.func_list = [{'dw_at_name': 'void', 'dw_at_type': None}]
		else:
		
			for func_param in self.func_param_range:
				#print(func_param)
				if func_param.find('<3>')>=0:
					continue
				dw_at_name = re.findall(r'DW_AT_name.*: (.+)\n', func_param, re.I)[0]
				dw_at_type = '><' + re.findall(r'DW_AT_type.*: <0x(.+)\n', func_param, re.I)[0]
				dict = {'dw_at_name': dw_at_name, 'dw_at_type': dw_at_type}
				self.func_list.append(dict)
		#print(self.func_list)
		
		return self.func_list
		
		
		
	def get_type_of_func(self, func):
		if self.no_parameter == True:
			func_header_range = re.findall(r'^.*?<1>', self.func_range, re.I|re.S)[0]
		else:
			func_header_range = re.findall(r'^.*?<2>', self.func_range, re.I|re.S)[0]

		type = re.findall(r'DW_AT_type.*: <0x(.+)\n', func_header_range, re.I)
		if len(type) == 0:
			return None
		else:
			location = '><' + type[0]
			return location
			
	def _merge_type_name_parameter(self, pre_value, func):
		param_value, subroutine_value = self.get_value_at_location(func['dw_at_type'])
		if param_value[-1] == '*':
			value = pre_value + param_value + func['dw_at_name'] + ', '
		elif param_value[-1] == ')':
			#print(pre_value)
			#print(param_value)
			#print(func['dw_at_name'])
			#print(subroutine_value)
			value = pre_value + param_value[:-1] + func['dw_at_name'] + ')' + subroutine_value + ', '
		else:
			value = pre_value + param_value + ' ' + func['dw_at_name'] + ', '
		
		return value
		
	def get_api_from_EFI(self, func):
		#print('* ' + func)
		# flag to judge whether the function contains parameter
		self.no_parameter = False 
		result = self.get_types_and_names_for_func(func)
		if result == -1:
			return -1
		func_type_location = self.get_type_of_func(func)
		#print(func_type_location)
		
		
		# get the function type
		if func_type_location is None:
			func_type_value = 'void '
			subroutine_value = -1
		else:
			func_type_value, subroutine_value = self.get_value_at_location(func_type_location)
			if subroutine_value == -1:
				if func_type_value[-1] != '*':
					func_type_value = func_type_value + ' '
			else:
				func_type_value = func_type_value[:-1]
		
		# merge the function type and name
		value = func_type_value + self.func_name + '('
		
		# merge the function type, name and parameters 
		if len(self.func_list) == 0:
			value = value + 'void)'
		elif len(self.func_list) == 1 and self.func_list[0]['dw_at_name'] == 'void':
			value =value + 'void)'
		else:
			for func in self.func_list:
				value = self._merge_type_name_parameter(value, func)
			value = value[:-2] + ')'
			if subroutine_value != -1:
				value = value + ') ' + subroutine_value
		print(value)
		return value

if __name__ == '__main__':
	path = sys.argv[1]
	func = sys.argv[2]
	elf = Read_ELF(path)
	elf.get_api_from_EFI(func)