Added support for one callback. Will be used by the saml2test tool.

1 files changed, 8 insertions, 2 deletions

diff --git a/src/saml2/entity.py b/src/saml2/entity.py
index 174d2e83..2f1c4d52 100644
--- a/src/saml2/entity.py
+++ b/src/saml2/entity.py
@@ -124,7 +124,7 @@ def create_artifact(entity_id, message_handle, endpoint_index=0):
 
 class Entity(HTTPBase):
     def __init__(self, entity_type, config=None, config_file="",
-                 virtual_organization=""):
+                 virtual_organization="", msg_cb=None):
         self.entity_type = entity_type
         self.users = None
 
@@ -177,6 +177,8 @@ class Entity(HTTPBase):
         else:
             self.sourceid = {}
 
+        self.msg_cb = msg_cb
+
     def _issuer(self, entityid=None):
         """ Return an Issuer instance """
         if entityid:
@@ -465,7 +467,6 @@ class Entity(HTTPBase):
                 kwargs[key] = val
 
         req = request_cls(**kwargs)
-        reqid = req.id
 
         if destination:
             req.destination = destination
@@ -479,6 +480,11 @@ class Entity(HTTPBase):
         if nsprefix:
             req.register_prefix(nsprefix)
 
+        if self.msg_cb:
+            req = self.msg_cb(req)
+
+        reqid = req.id
+
         if sign:
             return reqid, self.sign(req, sign_prepare=sign_prepare,
                                     sign_alg=sign_alg, digest_alg=digest_alg)