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

import argparse
import boto3, time, json

parser = argparse.ArgumentParser()

print('\n==============================================================')
print('traceLambda - Lambda Function Tracing Script (beta)')
print('(c) Blueinfy solutions pvt. ltd.')
print('==============================================================\n')


parser.add_argument('-i', action='store', dest='id',
                    help='use id for trace fetching')

parser.add_argument('-f', action='store', dest='function',
                    help='use function for trace fetching')

parser.add_argument('-x', action='store_true', default=False,
                    dest='xray',
                    help='use xRay to get trace of 120 seconds')



results = parser.parse_args()

if(results.id==None and results.function==None and results.xray==False):
	print(parser.parse_args(['-h']))


client = boto3.client('logs')
groups = client.describe_log_groups()
list_groups = groups["logGroups"]
getloggroups = []

for j in list_groups:
	getloggroups.append(j["logGroupName"])

if(results.id!=None):
	try:
		filter='\"'+results.id+'\"'
		for k in getloggroups:
			print("(+)Searching logs in "+k)
			res = client.filter_log_events(logGroupName=k, filterPattern=filter)
			all_events = []
			all_events = res["events"]
			for i in all_events:
				print("-------------------------------------------------")
				print('    (-)'+i["message"])
	except Exception as error:
		print('    (-)Error:',error)

if(results.function!=None):
	try:
		for k in getloggroups:
			if(k.find(results.function)>0):
				print("(+)Dumping traces from "+k)
				res = client.filter_log_events(logGroupName=k)
				all_events = []
				all_events = res["events"]
				
				for i in all_events:
					print("-------------------------------------------------")
					print('    (-)'+i["message"])
	except Exception as error:
		print('    (-)Error:',error)

if(results.xray==True):
	try:
		client = boto3.client('xray')

		response = client.get_trace_summaries(
		    StartTime=int(time.time())-120,
		    EndTime=int(time.time()),
		)

		traces = response["TraceSummaries"]
		for i in traces:
		        print("-------------------------------------------------")
		        print(i["Id"])
		        full_trace = client.batch_get_traces(TraceIds=[i["Id"]])
		        print(json.dumps(full_trace,indent=3))
		        print("-------------------------------------------------")
	except Exception as error:
		print('    (-)Error:',error)
