From 68344ab119d5f358bfd7639361871e692c4cda4a Mon Sep 17 00:00:00 2001
From: Clint Byrum <clint@fewbar.com>
Date: Wed, 27 May 2015 11:45:59 -0700
Subject: [PATCH] Fix artifact code for python3

Strings/bytes issues abound when hashing/encoding things.
---
 src/saml2/entity.py       | 20 +++++++++++++-------
 tests/test_64_artifact.py | 12 ++++++------
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/src/saml2/entity.py b/src/saml2/entity.py
index 6d868b5..60de964 100644
--- a/src/saml2/entity.py
+++ b/src/saml2/entity.py
@@ -34,7 +34,7 @@ from saml2.time_util import instant
 from saml2.s_utils import sid
 from saml2.s_utils import UnravelError
 from saml2.s_utils import error_status_factory
-from saml2.s_utils import rndstr
+from saml2.s_utils import rndbytes
 from saml2.s_utils import success_status_factory
 from saml2.s_utils import decode_base64_and_inflate
 from saml2.s_utils import UnsupportedBinding
@@ -73,7 +73,7 @@ logger = logging.getLogger(__name__)
 
 __author__ = 'rolandh'
 
-ARTIFACT_TYPECODE = '\x00\x04'
+ARTIFACT_TYPECODE = b'\x00\x04'
 
 SERVICE2MESSAGE = {
     "single_sign_on_service": AuthnRequest,
@@ -103,11 +103,17 @@ def create_artifact(entity_id, message_handle, endpoint_index=0):
     :param endpoint_index:
     :return:
     """
+    if not isinstance(entity_id, six.binary_type):
+        entity_id = entity_id.encode('utf-8')
     sourceid = sha1(entity_id)
 
-    ter = "%s%.2x%s%s" % (ARTIFACT_TYPECODE, endpoint_index,
-                          sourceid.digest(), message_handle)
-    return base64.b64encode(ter)
+    if not isinstance(message_handle, six.binary_type):
+        message_handle = message_handle.encode('utf-8')
+    ter = b"".join((ARTIFACT_TYPECODE, 
+                    ("%.2x" % endpoint_index).encode('ascii'),
+                    sourceid.digest(),
+                    message_handle))
+    return base64.b64encode(ter).decode('ascii')
 
 
 class Entity(HTTPBase):
@@ -1115,8 +1121,8 @@ class Entity(HTTPBase):
         :param endpoint_index:
         :return:
         """
-        message_handle = sha1("%s" % message)
-        message_handle.update(rndstr())
+        message_handle = sha1(str(message).encode('utf-8'))
+        message_handle.update(rndbytes())
         mhd = message_handle.digest()
         saml_art = create_artifact(self.config.entityid, mhd, endpoint_index)
         self.artifact[saml_art] = message
diff --git a/tests/test_64_artifact.py b/tests/test_64_artifact.py
index 52a6096..8b3bb26 100644
--- a/tests/test_64_artifact.py
+++ b/tests/test_64_artifact.py
@@ -54,14 +54,14 @@ def get_msg(hinfo, binding, response=False):
 
 def test_create_artifact():
     b64art = create_artifact("http://sp.example.com/saml.xml",
-                             "aabbccddeeffgghhiijj")
+                             b"aabbccddeeffgghhiijj")
 
-    art = base64.b64decode(b64art)
+    art = base64.b64decode(b64art.encode('ascii'))
 
-    assert art[:2] == '\x00\x04'
+    assert art[:2] == ARTIFACT_TYPECODE
     assert int(art[2:4]) == 0
 
-    s = sha1("http://sp.example.com/saml.xml")
+    s = sha1(b"http://sp.example.com/saml.xml")
     assert art[4:24] == s.digest()
 
 SP = 'urn:mace:example.com:saml:roland:sp'
@@ -74,7 +74,7 @@ def test_create_artifact_resolve():
     #assert artifact[:2] == '\x00\x04'
     #assert int(artifact[2:4]) == 0
     #
-    s = sha1(SP)
+    s = sha1(SP.encode('ascii'))
     assert artifact[4:24] == s.digest()
 
     with closing(Server(config_file="idp_all_conf")) as idp:
@@ -116,7 +116,7 @@ def test_artifact_flow():
                                                [BINDING_HTTP_ARTIFACT],
                                                entity_id=idp.config.entityid)
 
-        hinfo = sp.apply_binding(binding, "%s" % artifact, destination, relay_state)
+        hinfo = sp.apply_binding(binding, artifact, destination, relay_state)
 
         # ========== @IDP ============