# AUTOGENERATED! DO NOT EDIT! File to edit: cognito.ipynb (unless otherwise specified).

__all__ = ['Cognito', 'addUser', 'credenFromIdToken', 'login', 'unAuthCreden', 'getAttribute', 'loginWithFb',
           'addVillaEmployee', 'confirm', 'update', 'refreshToken', 'getSecretHash', 'disableUser', 'enableUser',
           'setNewPassword', 'updateAttribute', 'deleteUser']

# Cell
import ujson as json
from pycognito import Cognito as CognitoClient
from pycognito.aws_srp import AWSSRP
from nicHelper.wrappers import add_method
from nicHelper.exception import errorString
from nicHelper.dictUtil import printDict, filterDt
from botocore.exceptions import ClientError

import boto3, botocore, base64, hmac, hashlib


# Cell
class Cognito:
  '''sdk to interact with cognito as a client'''
  def __init__(
    self,
    identityPoolId,
    userPoolId,
    accountId,
    clientId,
    clientSecret,
    region = 'ap-southeast-1'
    ):
    self.identityPoolId = identityPoolId
    self.userPoolId = userPoolId
    self.accountId = accountId
    self.client = boto3.client('cognito-idp', region_name=region)
    self.identityClient = boto3.client('cognito-identity', region_name=region)
    self.clientId = clientId
    self.clientSecret = clientSecret
    self.region = region

  def getCognitoClient(self, user, accessToken = None, idToken = None, refreshToken = None):
    u = CognitoClient(
        self.userPoolId,
        self.clientId,
        client_secret=self.clientSecret,
        username= user,
        access_token = accessToken,
        id_token = idToken
      )
    return u

  def getSrp(self, user, pw):
    return AWSSRP(
        username = user,
        password= pw,
        pool_id = self.userPoolId,
        client_id = self.clientId,
        client_secret= self.clientSecret ,
        client = self.client
      )



# Cell
import botocore
@add_method(Cognito)
def addUser(self,user, phone, pw, name, **kwargs):
  try:
    u = CognitoClient(
      self.userPoolId,
      self.clientId,
      client_secret=self.clientSecret,
      username= user
    )
    u.set_base_attributes(phone_number= phone, name=name )
    u.register(user , pw)
    return {'success': True}
  except boto3.client('cognito-idp').exceptions.UsernameExistsException as e:
    return {'success': False, 'error': 'user exists'}
  except Exception as e:
    return {'success': False, 'error': e , 'traceback': errorString()}


# Cell
@add_method(Cognito)
def credenFromIdToken(self, logins):
  '''get aws credentials from cognito token'''
  ## get identityId
  identityResponse = self.identityClient.get_id(
      AccountId=self.accountId,
      IdentityPoolId=self.identityPoolId,
      Logins = logins
    )
  identityId = identityResponse['IdentityId']
  ## get aws credentials
  credentials_response = self.identityClient.get_credentials_for_identity(
      IdentityId=identityId,
      Logins = logins
      )
  return credentials_response['Credentials']

# Cell
@add_method(Cognito)
def login(self, user:str , pw:str, *args, **kwargs ):
  '''login to pool using cognito'''

  srp = self.getSrp(user,pw)
  try:
    tokens = srp.authenticate_user()
  except ClientError as e:
    print(type(e))
    if e.response['Error']['Code'] == 'NotAuthorizedException':
      return False, 'wrong password'
    return False, 'other problems'

  ##########################################################
  # extract accessToken etc
  accessToken = tokens['AuthenticationResult']['AccessToken']
  refreshToken = tokens['AuthenticationResult']['RefreshToken']
  idToken = tokens['AuthenticationResult']['IdToken']

  ####### Getting the IAM credentials############################
  logins = {f'cognito-idp.ap-southeast-1.amazonaws.com/{self.userPoolId}': idToken}
#   printDict(logins)
  creden = self.credenFromIdToken(logins)
  creden.update(tokens['AuthenticationResult'])
  creden['Expiration'] = creden['Expiration'].timestamp()
  ############ get user attributes
  userResult = self.client.admin_get_user(UserPoolId=self.userPoolId,Username=user)
  userAttributes = userResult.pop('UserAttributes')
  kwUserAttrib = {attrib['Name']:attrib['Value'] for attrib in userAttributes}
#   print(kwUserAttrib)
  creden['userInfo'] = userResult
  creden['kwUserAttrib'] = kwUserAttrib

  return True, filterDt(creden)




# Cell
@add_method(Cognito)
def unAuthCreden(self):
  identityResponse = self.identityClient.get_id(
    AccountId=self.accountId,
    IdentityPoolId=self.identityPoolId
    )
  identityId = identityResponse['IdentityId']
  credentialResponse = self.identityClient.get_credentials_for_identity(
      IdentityId=identityId
      )
  return credentialResponse['Credentials']

# Cell
@add_method(Cognito)
def getAttribute(self, user):
  userResult = self.client.admin_get_user(UserPoolId=self.userPoolId,Username=user)
  try: userResult['UserAttributes'] = {attrib['Name']:attrib['Value'] for attrib in userResult['UserAttributes']}
  except Exception as e: raise Exception(f'error parsing attributes from {userResult}\n{e}')
  return filterDt(userResult)

# Cell
@add_method(Cognito)
def loginWithFb(self, fbKey):
  logins = {'graph.facebook.com':fbKey}
  printDict(logins)
  creden = self.credenFromIdToken(logins)
  creden['Expiration'] = creden['Expiration'].timestamp()
  return creden


# Cell
@add_method(Cognito)
def addVillaEmployee(self,user, cashierCode, name, phone, pw, email = ""):
  try:
    u = CognitoClient(
      self.userPoolId,
      self.clientId,
      client_secret=self.clientSecret,
      username= user
    )
    u.set_base_attributes(phone_number= phone, name=name, address=email)
    u.add_custom_attributes(cashierCode=cashierCode)
    u.register(user , pw)
    return {'success': True}
  except boto3.client('cognito-idp').exceptions.UsernameExistsException as e:
    return {'success': False, 'error': 'user exists'}
  except Exception as e:
    return {'success': False, 'error': e , 'traceback': errorString()}


# Cell
@add_method(Cognito)
def confirm(self, user, code):
  u = CognitoClient(
    self.userPoolId,
    self.clientId,
    client_secret=self.clientSecret,
    username= user
  )
  try:
    result = u.confirm_sign_up(code ,username=user)
    return {'success': True}
  except Exception as e:
    return {'success': False, 'error': e , 'traceback': errorString()}


# Cell
@add_method(Cognito)
def update(user):

  u = Cognito('your-user-pool-id','your-client-id',
    id_token='id-token',refresh_token='refresh-token',
    access_token='access-token')

  u.send_verification(attribute='email')
  return {'update': 'update'}


# Cell
@add_method(Cognito)
def refreshToken(self, user,refreshToken):
    try:
        client = boto3.client('cognito-idp')
#         print(refreshToken)
        return client.initiate_auth(
            ClientId=self.clientId,
            AuthFlow='REFRESH_TOKEN_AUTH',
            AuthParameters={
                'REFRESH_TOKEN': refreshToken,
                'SECRET_HASH': self.getSecretHash(user)
            }
        )
    except botocore.exceptions.ClientError as e:
        return e.response

# Cell
@add_method(Cognito)
def getSecretHash(self, username):
      message = username + self.clientId
      dig = hmac.new(self.clientSecret.encode(), msg=message.encode('UTF-8'),
                     digestmod=hashlib.sha256).digest()
      return base64.b64encode(dig).decode()


# Cell
@add_method(Cognito)
def disableUser(self, username):
  result = self.client.admin_disable_user(
    UserPoolId=self.userPoolId,
    Username=username
  )
  return result

# Cell
@add_method(Cognito)
def enableUser(self, username):
  result = self.client.admin_enable_user(
    UserPoolId=self.userPoolId,
    Username=username
  )
  return result

# Cell
@add_method(Cognito)
def setNewPassword(self, user, newPassword):
  try:
    response = self.client.admin_set_user_password(
      UserPoolId=self.userPoolId,
      Username=user,
      Password=newPassword,
      Permanent=True
    )
  except:
    return False, errorString()
  return True, response

# Cell
@add_method(Cognito)
def updateAttribute(self,user:str, attributes:dict):
  listAttributes = [ {'Name': k, 'Value': v} for k,v in attributes.items()]
  try:
    self.client.admin_update_user_attributes(
      UserPoolId=self.userPoolId,
      Username=user,
      UserAttributes=listAttributes,
    )
    response = self.getAttribute('nic1')
  except Exception as e:
    return False, f'error updating {e}'
  return True, response

# Cell
@add_method(Cognito)
def deleteUser(self,user:str):
  response = self.client.admin_delete_user(
    UserPoolId=self.userPoolId,
    Username=user
  )
  return response