'''
tunnelLambda - 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, base64
from http.server import BaseHTTPRequestHandler, HTTPServer

parser = argparse.ArgumentParser()

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



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

parser.add_argument('-p', action='store', dest='port',
                    help='supply port for listening')

results = parser.parse_args()
if(results.function==None or results.port==None):
	print(parser.parse_args(['-h']))

FUNCTION = results.function
PORT = int(results.port)


class tunnel(BaseHTTPRequestHandler):
    def _set_response(self):
        self.send_response(200)
        self.send_header('Content-type', 'application/json')
        self.end_headers()

    def do_GET(self):
        #self._set_response()
        self.send_response(200)
        self.send_header('Content-type', 'text/html')
        self.end_headers()
        htmlfile = open("./tunnel-config/get.html","r").read()
        htmlwithport = htmlfile.replace("$port$",results.port)
        finalhtml = htmlwithport.replace("$function$",results.function)
        self.wfile.write(finalhtml.encode('utf-8'))

    def do_POST(self):
        content_length = int(self.headers['Content-Length'])
        eventstream = self.rfile.read(content_length)
        try:
        	client = boto3.client('lambda')
        	res = client.invoke(FunctionName=FUNCTION,LogType='Tail', Payload=eventstream)
        	response = res['Payload'].read().decode('utf-8')
        	self._set_response()
        	self.wfile.write(response.encode('utf-8'))
        except Exception as e:
        	self._set_response()
        	self.wfile.write("Error : check exception on shell".encode('utf-8'))
        	print(e)
       
def run(server_class=HTTPServer, handler_class=tunnel, port=PORT):
    
    server_address = ('', port)
    httpd = server_class(server_address, handler_class)
    print("Established tunnel for lambda function - ",FUNCTION)
    print("Listening on Port - ",PORT)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass
    httpd.server_close()

run()


   


