'''
scanLambda - Lambda Function Scanner Script
(c) Blueinfy solutions pvt. ltd.
This is beta version of the toolkit and it is in the development
phase.
==============================================================
'''

import argparse, boto3, sys, json, base64, re

parser = argparse.ArgumentParser()

print('\n==============================================================')
print('scanLambda - Lambda Scanner Script (beta)')
print('(c) Blueinfy solutions pvt. ltd.')
print('==============================================================\n')


parser.add_argument('-i', action='store_true', dest='invoke', default=False,
                    help='use for invoke the function')

parser.add_argument('-f', action='store', dest='function',
                    help='use function for scan/invoke')

parser.add_argument('-s', action='store_true', dest='scan', default=False,
                    help='use for scan the function')

parser.add_argument('-e', action='store', dest='event',
                    help='supply event file')
parser.add_argument('-p', action='store_true', default=False,
                    dest='scanperm',
                    help='Scan for permission issues')

results = parser.parse_args()
client = boto3.client('lambda')
iam = boto3.client('iam')

def getFunctions():
	all_functions = raw_functions["Functions"]
	list_of_functions = []
	for i in all_functions:
		list_of_functions.append(i["FunctionName"])
	return list_of_functions

if(results.invoke!=False and results.function!=None and results.event!=None):

	print("(+)Configuring Invoking ...")
	print('	(-) Loading event from file ...')
	eventfile = results.event
	event = open(eventfile,"r")
	eventstring = event.read()
	eventstream = json.dumps(json.loads(eventstring))

	res = client.invoke(FunctionName=results.function,LogType='Tail', Payload=eventstream)
	print("	(-) Request Id ==> "+res["ResponseMetadata"]["RequestId"])
	response = res['Payload'].read().decode('utf-8')
	print("	(-) Response ==> ")
	print("	(-) " + response)
	print("	(-) Log ==>" + base64.b64decode(res['LogResult']).decode('utf-8'))

elif(results.scan!=False and results.function!=None and results.event!=None):

	print("(+)Configuring scanner ...")
	print('	(-)Loading payload from file ...')
	payloadfile = "./scan-config/payload.txt"
	payload = open(payloadfile,"r")
	temp = payload.read()
	injections = temp.split("\n")

	print('	(-)Loading regex from file ...')
	regexfile = "./scan-config/regex.txt"
	regexlist = open(regexfile,"r")
	temp = regexlist.read()
	regexsplits = temp.split("\n")

	print('	(-)Loading event from file ...')
	eventfile = results.event
	event = open(eventfile,"r")
	eventstring = event.read()

	for i in injections:
		print("(+) ==> ----------------------------")
		print("	(-) Injecting payload - " + i)
		final = eventstring.replace("$fuzz$",i)
		eventstream = json.dumps(json.loads(final))

		res = client.invoke(FunctionName=results.function,LogType='Tail', Payload=eventstream)
		print("	(-) Request Id ==> "+res["ResponseMetadata"]["RequestId"])
		response = res['Payload'].read().decode('utf-8')
		print("	(-) Response Analysis ")
		print(response)
		
		for j in regexsplits:
			m = re.search(j,response,re.IGNORECASE)
			if m:
				print("	(-) Regex found: "+j)
				print("	(-) " + response)
				print("	(-) Log ==>" + base64.b64decode(res['LogResult']).decode('utf-8'))

elif(results.scanperm!=False):
	raw_functions = client.list_functions()
	function_names = getFunctions()
	print('(+) Loading permission rules from file ...')
	permfile = "./scan-config/permission.txt"
	perm = open(permfile,"r")
	temp = perm.read()
	permlist = temp.split("\n")
	print("(+) Scanning Lambda functions Permissions...")
	for j in function_names:
		print("  (-)"+j)
		myfunc = client.get_function(FunctionName=j)
		temp = myfunc["Configuration"]
		rolelist = temp["Role"].split('/')
		length = len(rolelist)
		policy = iam.list_attached_role_policies(RoleName=temp["Role"].split('/')[length-1])
		attachedpolicy = policy["AttachedPolicies"]
		for i in attachedpolicy:			
			rolepolicy = iam.get_policy(PolicyArn=i["PolicyArn"])
			getpolicyversion = iam.get_policy_version(PolicyArn=i["PolicyArn"], VersionId=rolepolicy['Policy']['DefaultVersionId'])
			perm = json.dumps(getpolicyversion['PolicyVersion']['Document']['Statement'])
			print("		==> "+perm)
			for i in permlist:
				m = re.search(i,perm,re.IGNORECASE)
				if m:
					print("		(-) Permission issue found: "+i+"\n")
		print("\n")
		
else:

	print(parser.parse_args(['-h']))
