import base64, hashlib, jwt, os, time
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad
from dotenv import load_dotenv
from lib.commonlib import CommonLib

load_dotenv()
secret = bytes(os.getenv('JWT_SECRET'), 'utf-8')
CommonLib = CommonLib()


class LockerJwt:
    def __init__(self, request):
        try:
            bearer = request.headers.get('authorization', None)
            jwt_token = bearer.split(" ")[-1] if bearer is not None else None
            token = jwt_token if jwt_token is not None else request.headers.get('jtoken')
            device_security_id = request.headers.get('device-security-id')
            self.jwt_token = token
            self.device_security_id = device_security_id
        except Exception as e:
            return 401, {"status": "error", "error_description": 'Exception:LockerJwt:__init__:: ' + str(e)}

    @staticmethod
    def aes_decryption(filtered_cipher_text, secret_key):
        try:
            iv = bytes(16 * '\x00', 'utf-8')
            encode_cipher = base64.b64decode(filtered_cipher_text)
            aes_obj = AES.new(secret_key, AES.MODE_CBC, iv)
            return unpad(aes_obj.decrypt(encode_cipher), AES.block_size).decode('utf-8')
        except Exception as e:
            return 401, {"status": "error", "error_description": 'Exception:LockerJwt:aes_decryption:: ' + str(e)}

    def _pad(self, s):
        bs = AES.block_size
        return s + (bs - len(s) % bs) * chr(bs - len(s) % bs)
    
    def aes_encryption(self, raw,secret_key):
        raw = self._pad(raw)
        iv = bytes(16 * '\x00', 'utf-8')
        cipher = AES.new(secret_key, AES.MODE_CBC, iv)
        return base64.b64encode(cipher.encrypt(raw.encode())).decode('utf-8').replace('+', '---')


    def jwt_login(self):
        try:
            
            if not bool(int(os.getenv('JWT_ENABLED'))):
                return 200,{'digilockerid':os.getenv('JWT_LOCKERID'),'device_security_id':self.device_security_id,'token':self.jwt_token, 'username':os.getenv('JWT_LOCKERID')}
            token = self.jwt_token
            if token is None:
                return 401, {"status": "error", "error_description": "Jtoken is missing."}
                
            jwt_res = jwt.decode(token, os.getenv('JWT_SECRET'), audience='DIGILOCKER', algorithms=['HS256'])
            data = jwt_res.get('data')
            enc_username = data.get('username')
            username = self.aes_decryption(enc_username.replace('---', '+'), secret)
            digilockerid = data.get('digilockerid')
            did_sign = None
            did = None
            if data.get('didsign'):
                did_sign_enc = data.get('didsign')
                did_sign = self.aes_decryption(did_sign_enc.replace('---', '+'), secret)
            device_security_id = self.device_security_id
            # device_security_id is mandatory
            if device_security_id:
                result = hashlib.md5(device_security_id.encode())
                did = result.hexdigest()
            
            else:
                return 401, {"status": "error", "error_description": "Unauthorised Access!"}
            
            if did_sign and (did is None or did_sign.find(did) == -1):
                return 401, {"status": "error", "error_description": "Unauthorised Access!"}
            
            return 200, {'digilockerid':digilockerid,'device_security_id': device_security_id,'token': token, 'username':username}
        except Exception as e:
            return 401, {"status": "error", "error_description": 'Exception:LockerJwt:jwt_login:: ' + str(e)}
    def jwt_generate(self,digilockerid):
        try:
            secret = bytes(os.getenv('JWT_SECRET'), 'utf-8')
            ts= int(time.time())
            payload ={
                "iss":"https://digilocker.gov.in",
                "aud":"DIGILOCKER",
                "iat":ts,
                "nbf":ts-60,
                "exp":ts+84000,
                "data":{"status":"success",
                "username":self.aes_encryption(digilockerid,secret),
                "digilockerid":digilockerid,
                "didsign":self.aes_encryption(self.device_security_id,secret)}
            }
            encoded = jwt.encode(payload, os.getenv('JWT_SECRET'), algorithm="HS256")
            return encoded
        except Exception as e:
            return 401, {"status": "error", "error_description": 'Exception:LockerJwt:jwt_login:: ' + str(e)}