Refactor handling of nameid format

Signed-off-by: Ivan Kanakarakis <[email protected]>

Author: c00kiemon5ter

src/satosa/frontends/saml2.py

@@ -16,9 +16,11 @@
 from saml2.config import IdPConfig
 from saml2.extension.ui import NAMESPACE as UI_NAMESPACE
 from saml2.metadata import create_metadata_string
-from saml2.saml import NameID, NAMEID_FORMAT_TRANSIENT, \
-    NAMEID_FORMAT_PERSISTENT, NAMEID_FORMAT_EMAILADDRESS, \
-    NAMEID_FORMAT_UNSPECIFIED
+from saml2.saml import NameID
+from saml2.saml import NAMEID_FORMAT_TRANSIENT
+from saml2.saml import NAMEID_FORMAT_PERSISTENT
+from saml2.saml import NAMEID_FORMAT_EMAILADDRESS
+from saml2.saml import NAMEID_FORMAT_UNSPECIFIED
 from saml2.samlp import name_id_policy_from_string
 from saml2.server import Server
 
@@ -45,14 +47,16 @@ def saml_name_id_format_to_hash_type(name_format):
     :param name_format: SAML2 name format
     :return: satosa format
     """
-    if name_format == NAMEID_FORMAT_PERSISTENT:
-        return UserIdHashType.persistent
-    elif name_format == NAMEID_FORMAT_EMAILADDRESS:
-        return UserIdHashType.emailaddress
-    elif name_format == NAMEID_FORMAT_UNSPECIFIED:
-        return UserIdHashType.unspecified
+    name_id_format_to_hash_type = {
+        NAMEID_FORMAT_TRANSIENT:    UserIdHashType.transient,
+        NAMEID_FORMAT_PERSISTENT:   UserIdHashType.persistent,
+        NAMEID_FORMAT_EMAILADDRESS: UserIdHashType.emailaddress,
+        NAMEID_FORMAT_UNSPECIFIED:  UserIdHashType.unspecified,
+    }
 
-    return UserIdHashType.transient
+    return name_id_format_to_hash_type.get(
+            name_format,
+            UserIdHashType.transient)
 
 
 def hash_type_to_saml_name_id_format(hash_type):
@@ -64,16 +68,16 @@ def hash_type_to_saml_name_id_format(hash_type):
     :param hash_type: satosa format
     :return: pySAML2 name format
     """
-    if hash_type is UserIdHashType.transient:
-        return NAMEID_FORMAT_TRANSIENT
-    elif hash_type is UserIdHashType.persistent:
-        return NAMEID_FORMAT_PERSISTENT
-    elif hash_type is UserIdHashType.emailaddress:
-        return NAMEID_FORMAT_EMAILADDRESS
-    elif hash_type is UserIdHashType.unspecified:
-        return NAMEID_FORMAT_UNSPECIFIED
-
-    return NAMEID_FORMAT_PERSISTENT
+    hash_type_to_name_id_format = {
+        UserIdHashType.transient:    NAMEID_FORMAT_TRANSIENT,
+        UserIdHashType.persistent:   NAMEID_FORMAT_PERSISTENT,
+        UserIdHashType.emailaddress: NAMEID_FORMAT_EMAILADDRESS,
+        UserIdHashType.unspecified:  NAMEID_FORMAT_UNSPECIFIED,
+    }
+
+    return hash_type_to_name_id_format.get(
+            hash_type,
+            NAMEID_FORMAT_PERSISTENT)
 
 
 class SAMLFrontend(FrontendModule, SAMLBaseModule):

src/satosa/internal_data.py

@@ -6,6 +6,7 @@
 import hashlib
 from enum import Enum
 
+
 class UserIdHashType(Enum):
     """
     All different user id hash types
@@ -56,7 +57,9 @@ def hash_data(salt, value):
         :param value: value to hash together with the salt
         :return: hash value (SHA512)
         """
-        return hashlib.sha512((value + salt).encode("utf-8")).hexdigest()
+        data = '{value}{salt}'.format(value=value, salt=salt).encode()
+        hash = hashlib.sha512(data).hexdigest()
+        return hash
 
     @staticmethod
     def hash_type(state):
@@ -82,23 +85,36 @@ def hash_id(salt, user_id, requester, state):
         :param state: The current state
         :return: the internal_response containing the hashed user ID
         """
+        hash_type_to_format = {
+            UserIdHashType.transient:    '{id}{req}{time}',
+            UserIdHashType.persistent:   '{id}{req}',
+            UserIdHashType.pairwise:     '{id}{req}',
+            UserIdHashType.public:       '{id}',
+            UserIdHashType.emailaddress: '{id}',
+            UserIdHashType.unspecified:  '{id}',
+        }
+
+        format_args = {
+            'id': user_id,
+            'req': requester,
+            'time': datetime.datetime.utcnow().timestamp(),
+        }
+
         hash_type = UserIdHasher.hash_type(state)
-        if hash_type == UserIdHashType.transient:
-            timestamp = datetime.datetime.now().time()
-            user_id = "{req}{time}{id}".format(req=requester, time=timestamp,
-                                               id=user_id)
-        elif (hash_type == UserIdHashType.persistent or
-              hash_type == UserIdHashType.pairwise):
-            user_id = "{req}{id}".format(req=requester, id=user_id)
-        elif hash_type == UserIdHashType.public:
-            user_id = "{id}".format(id=user_id)
-        elif (hash_type == UserIdHashType.emailaddress or
-              hash_type == UserIdHashType.unspecified):
-            return user_id
+        try:
+            fmt = hash_type_to_format[hash_type]
+        except KeyError as e:
+            raise ValueError('Unknown hash type: {}'.format(hash_type)) from e
         else:
-            raise ValueError("Unknown hash type: '{}'".format(hash_type))
-
-        return UserIdHasher.hash_data(salt, user_id)
+            user_id = fmt.format(**format_args)
+
+        hasher = (
+            (lambda salt, value: value)
+            if hash_type in [
+                UserIdHashType.emailaddress,
+            ]
+            else UserIdHasher.hash_data)
+        return hasher(salt, user_id)
 
 
 class AuthenticationInformation(object):

tests/satosa/frontends/test_saml2.py

@@ -14,9 +14,10 @@
 from saml2.entity_category.edugain import COCO
 from saml2.entity_category.refeds import RESEARCH_AND_SCHOLARSHIP
 from saml2.entity_category.swamid import SFS_1993_1153, RESEARCH_AND_EDUCATION, EU, HEI, NREN
-from saml2.saml import NameID, NAMEID_FORMAT_TRANSIENT, \
-    NAMEID_FORMAT_PERSISTENT, NAMEID_FORMAT_EMAILADDRESS, \
-    NAMEID_FORMAT_UNSPECIFIED
+from saml2.saml import NAMEID_FORMAT_TRANSIENT
+from saml2.saml import NAMEID_FORMAT_PERSISTENT
+from saml2.saml import NAMEID_FORMAT_EMAILADDRESS
+from saml2.saml import NAMEID_FORMAT_UNSPECIFIED
 from saml2.samlp import NameIDPolicy
 
 from satosa.attribute_mapping import AttributeMapper