ergo70 · February 25, 2022 20:59 · Mar 14, 2020 · Mar 13, 2020 · Mar 13, 2020 · Mar 12, 2020
diff --git a/cognito_PAM.py b/cognito_PAM.py
@@ -9,6 +9,7 @@
 
 _PAM_DUMMY_PROMPT = ''
 
+_PUBKEYS = None
 
 def _get_credentials(pamh):
     user = None
@@ -49,7 +50,10 @@ def _jwks2PEM(jwks_key):
 
 
 def _verify_claim(region, token, pool_id, client_id=None):
-    pubkeys = _get_public_keys_for_cognito_pool(region, pool_id)
+    pubkeys = _PUBKEYS
+
+    if not pubkeys:
+        pubkeys = _get_public_keys_for_cognito_pool(region, pool_id)
 
     claim_header = jwt.get_unverified_header(token)
 

diff --git a/cognito_PAM.py b/cognito_PAM.py
@@ -87,7 +87,7 @@ def _cognito_auth_handler(region, pool_id, client_id, user, password):
 
         return not (None in (id_claim, access_claim))
 
-    return True
+    return False
 
 
 def pam_sm_authenticate(pamh, flags, argv):

diff --git a/cognito_PAM.py b/cognito_PAM.py
@@ -1,30 +1,67 @@
 import boto3
+import jwt
+from json import dumps
+from requests import get
 from botocore import UNSIGNED
 from botocore.config import Config
 from warrant.aws_srp import AWSSRP
+from cryptography.hazmat.primitives import serialization
 
 _PAM_DUMMY_PROMPT = ''
 
 
 def _get_credentials(pamh):
-    _user = None
-    _password = None
+    user = None
+    password = None
 
     try:
-        _user = pamh.get_user(None)
+        user = pamh.get_user(None)
 
         if pamh.authtok == None:  # If authtok not set, start conversation to get it
-            _msg = pamh.Message(pamh.PAM_PROMPT_ECHO_OFF, _PAM_DUMMY_PROMPT)
-            _conv = pamh.conversation(_msg)
-            pamh.authtok = _conv.resp
+            msg = pamh.Message(pamh.PAM_PROMPT_ECHO_OFF, _PAM_DUMMY_PROMPT)
+            conv = pamh.conversation(msg)
+            pamh.authtok = conv.resp
 
     except pamh.exception as e:
-        # return e.pam_result
         return None
 
-    _password = pamh.authtok
+    password = pamh.authtok
 
-    return (_user, _password)
+    return user, password
+
+
+def _get_public_keys_for_cognito_pool(region, pool_id):
+    pubkeys = get('https://cognito-idp.{}.amazonaws.com/{}/.well-known/jwks.json'.format(
+        region, pool_id))
+
+    pubkeys = pubkeys.json()
+    pubkeys = pubkeys.get('keys')
+
+    return pubkeys
+
+
+def _jwks2PEM(jwks_key):
+    public_key = jwt.algorithms.RSAAlgorithm.from_jwk(dumps(jwks_key))
+    PEM_key = public_key.public_bytes(
+        encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo)
+
+    return PEM_key
+
+
+def _verify_claim(region, token, pool_id, client_id=None):
+    pubkeys = _get_public_keys_for_cognito_pool(region, pool_id)
+
+    claim_header = jwt.get_unverified_header(token)
+
+    for key in pubkeys:
+        if key.get('kid') == claim_header.get('kid'):
+            verification_algorithm = key.get('alg')
+            verification_key = _jwks2PEM(key)
+            claim = jwt.decode(token, verification_key,
+                               algorithm=verification_algorithm, audience=client_id)
+            return claim
+
+    return None
 
 
 def _cognito_auth_handler(region, pool_id, client_id, user, password):
@@ -40,20 +77,26 @@ def _cognito_auth_handler(region, pool_id, client_id, user, password):
     tokens = aws.authenticate_user()
 
     if tokens:
-        return True
+        tokens = tokens.get('AuthenticationResult')
+
+        id_token = tokens.get('IdToken')
+        access_token = tokens.get('AccessToken')
+
+        id_claim = _verify_claim(region, id_token, pool_id, client_id)
+        access_claim = _verify_claim(region, access_token, pool_id)
 
-    return False
+        return not (None in (id_claim, access_claim))
 
-# Authenticate user. Type: auth
+    return True
 
 
 def pam_sm_authenticate(pamh, flags, argv):
     if len(argv) != 4:
         return pamh.PAM_AUTH_ERR
 
-    _user, _password = _get_credentials(pamh)
+    user, password = _get_credentials(pamh)
 
-    if _cognito_auth_handler(argv[1], argv[2], argv[3], _user, _password):
+    if _cognito_auth_handler(argv[1], argv[2], argv[3], user, password):
         return pamh.PAM_SUCCESS
 
     return pamh.PAM_AUTH_ERR

diff --git a/cognito_PAM.py b/cognito_PAM.py
@@ -28,7 +28,7 @@ def _get_credentials(pamh):
 
 
 def _cognito_auth_handler(region, pool_id, client_id, user, password):
-    if None not in (region, pool_id, client_id, user, password):
+    if None in (region, pool_id, client_id, user, password):
         return False
 
     client = boto3.client('cognito-idp',

diff --git a/cognito_PAM.py b/cognito_PAM.py
@@ -0,0 +1,65 @@
+import boto3
+from botocore import UNSIGNED
+from botocore.config import Config
+from warrant.aws_srp import AWSSRP
+
+_PAM_DUMMY_PROMPT = ''
+
+
+def _get_credentials(pamh):
+    _user = None
+    _password = None
+
+    try:
+        _user = pamh.get_user(None)
+
+        if pamh.authtok == None:  # If authtok not set, start conversation to get it
+            _msg = pamh.Message(pamh.PAM_PROMPT_ECHO_OFF, _PAM_DUMMY_PROMPT)
+            _conv = pamh.conversation(_msg)
+            pamh.authtok = _conv.resp
+
+    except pamh.exception as e:
+        # return e.pam_result
+        return None
+
+    _password = pamh.authtok
+
+    return (_user, _password)
+
+
+def _cognito_auth_handler(region, pool_id, client_id, user, password):
+    if None not in (region, pool_id, client_id, user, password):
+        return False
+
+    client = boto3.client('cognito-idp',
+                          region_name=region, config=Config(signature_version=UNSIGNED))
+
+    aws = AWSSRP(username=user, password=password, pool_id=pool_id,
+                 client_id=client_id, client=client)
+
+    tokens = aws.authenticate_user()
+
+    if tokens:
+        return True
+
+    return False
+
+# Authenticate user. Type: auth
+
+
+def pam_sm_authenticate(pamh, flags, argv):
+    if len(argv) != 4:
+        return pamh.PAM_AUTH_ERR
+
+    _user, _password = _get_credentials(pamh)
+
+    if _cognito_auth_handler(argv[1], argv[2], argv[3], _user, _password):
+        return pamh.PAM_SUCCESS
+
+    return pamh.PAM_AUTH_ERR
+
+# Authorize user account. Type: account
+
+
+def pam_sm_acct_mgmt(pamh, flags, argv):
+    return pamh.PAM_SUCCESS
No results found