Untitled

import base64
from urllib.parse import urlencode
import http.cookies as cookies
import http.client 
import jwt 
from jwt import PyJWKClient

cognito_user_pool_id = "us-east-1_f1oN3uxu6"
cognito_client_id = "mi6m1p1j3c7q6hdfiapmar0cb"
cognito_client_secret = "1a877rudvu15nf90hm0v0j3rst3q4e9smhjlu9a16o603t291g1q"
    
def get_authorization_code(request) :
    credentials = f"{cognito_client_id}:{cognito_client_secret}"
    credentials_bytes = credentials.encode('utf-8')
    credentials_base64 = base64.b64encode(credentials_bytes).decode('utf-8')

    # print(credentials_base64)

    headers = {
        'content-type': 'application/x-www-form-urlencoded',
        'authorization': 'Basic ' + credentials_base64
    }


    code = request['querystring'].split("=")[1]
    data = {
        'grant_type': 'authorization_code',
        'redirect_uri': 'https://d3df1p0inqhb93.cloudfront.net/parseauth',
        'code': code
    }

    #print(code)

    url = "stylcoggoogle.auth.us-east-1.amazoncognito.com"
    conn = http.client.HTTPSConnection(url)
    conn.request("POST", "/oauth2/token", urlencode(data), headers)

    response = conn.getresponse()
    response_data = response.read().decode('utf-8')

    response_data = eval(response_data)
    
    print("RESPONSE_DATA THAT GET ID_TOKEN : ", response_data)

    my_cookie = cookies.SimpleCookie()
    my_cookie['token'] = response_data['id_token']
    my_cookie['token']['max-age'] = response_data['expires_in']
    my_cookie['token']['path'] = '/'
    my_cookie['token']['secure'] = True

    set_cookie_value = my_cookie.output(header = '')
    print("MY COOKIE VALUE : ", set_cookie_value)
    
    final_response = {
        'status': '302',
        'statusDescription': 'Found',
        'headers': {
            'location': [{
                'key': 'Location',
                'value': 'https://d3df1p0inqhb93.cloudfront.net/'
            }],
        'set-cookies' : [{
            'key': 'Set-Cookie',
            'value' : set_cookie_value
        }],
        'cache-control': [{ 
            'key': 'Cache-Control',
            'value': 'no-cache'
            }]
        }
    }

    print("SET COOKIES SUCCESS, RETURN FINAL_RESPONSE ", final_response)
    return final_response


def verify_jwt_token(token) :
    jwk_url = f"https://cognito-idp.us-east-1.amazonaws.com/{cognito_user_pool_id}/.well-known/jwks.json"

    try : 
        client = PyJWKClient(jwk_url)
        pub_key = client.get_signing_key_from_jwt(token).key
        validation = jwt.decode(token, pub_key, audience = cognito_client_id, algorithms = ["RS256"], options = {"verify_signature": True, "verify_exp": True})
        return True 
    
    except :
        return False 
    


def lambda_handler(event, context):
    # print("EVENT FROM CLOUDFRONT : ", event)
    
    request = event['Records'][0]['cf']['request']

    if (request['uri'].startswith('/parseauth')) :
        return get_authorization_code(request)

    if 'cookie' in request['headers'] : 
        cookie = request['headers']['cookie'][0]['value'].split("=")[1]
        print("COOKIE VALUE : ", cookie)

        if (1 != 1) : 
            print("SUCCESSFUL VALIDATION !!!")
            return request    
        
        print("FAIL VALIDATION !!! REDIRECT TO SIGN-IN PAGE AGAIN !!!")
    
    response = {
        'status': '302',
        'statusDescription': 'Found',
        'headers': {
            'location': [{
                'key': 'Location',
                'value': 'https://stylcoggoogle.auth.us-east-1.amazoncognito.com/login?client_id=mi6m1p1j3c7q6hdfiapmar0cb&response_type=code&scope=aws.cognito.signin.user.admin+email+openid+phone+profile&redirect_uri=https%3A%2F%2Fd3df1p0inqhb93.cloudfront.net%2Fparseauth'
                
                
            }]
        }
    }

    print("REDIRECT RESPONSE TO COGNITO SIGN-IN : ", response)
    return response