Added functionality needed by the saml2test tool.

Author: Roland Hedberg

src/saml2/argtree.py

@@ -0,0 +1,70 @@
+__author__ = 'roland'
+
+
+def find_paths(cls, arg, path=None, seen=None, res=None, lev=0):
+    if lev == 0 and res is None:
+        res = []
+
+    if path is None:
+        path = []
+
+    if seen is None:
+        seen = [cls]
+    else:
+        if cls in seen:
+            return None
+
+        seen.append(cls)
+
+    for cn, c in cls.c_children.values():
+        _path = path + [cn]
+        if cn == arg:
+            if res is not None:
+                res.append(_path)
+        else:
+            if isinstance(c, list):
+                _c = c[0]
+            else:
+                _c = c
+
+            find_paths(_c, arg, _path, seen, res)
+
+    for an, typ, mult in cls.c_attributes.values():
+        if an == arg:
+            if res is not None:
+                res.append(path + [an])
+
+    if lev == 0:
+        return res
+
+
+def set_arg(cls, arg, value):
+    res = []
+    for path in find_paths(cls, arg):
+        x = y = {}
+        for arc in path[:-1]:
+            y[arc] = {}
+            y = y[arc]
+        y[path[-1]] = value
+        res.append(x)
+
+    return res
+
+
+def add_path(tdict, path):
+    """
+
+    :param tdict: a dictionary representing a argument tree
+    :param path: a path list
+    :return: a dictionary
+    """
+    t = tdict
+    for step in path[:-2]:
+        try:
+            t = t[step]
+        except KeyError:
+            t[step] = {}
+            t = t[step]
+    t[path[-2]] = path[-1]
+
+    return tdict

tests/test_06_setarg.py

@@ -0,0 +1,44 @@
+from saml2 import saml
+from saml2.saml import Subject
+from saml2.samlp import Response
+from saml2.argtree import set_arg, add_path
+from saml2.argtree import find_paths
+
+__author__ = 'roland'
+
+
+def test_path():
+    result = find_paths(Subject, 'in_response_to')
+
+    assert result == [
+        ['subject_confirmation', 'subject_confirmation_data', 'in_response_to']]
+
+    result = find_paths(Response, 'in_response_to')
+
+    assert result == [
+        ['assertion', 'subject', 'subject_confirmation',
+         'subject_confirmation_data', 'in_response_to'],
+        ['in_response_to']
+    ]
+
+
+def test_set_arg():
+    r = set_arg(Subject, 'in_response_to', '123456')
+
+    assert r == [{'subject_confirmation': {
+        'subject_confirmation_data': {'in_response_to': '123456'}}}]
+
+
+def test_multi():
+    t = {}
+    t = add_path(t, ['subject_confirmation','method',saml.SCM_BEARER])
+    x = add_path(
+        t['subject_confirmation'],
+        ['subject_confirmation_data','in_response_to','1234'])
+
+    print(t)
+    assert t == {
+        'subject_confirmation': {
+            'subject_confirmation_data': {'in_response_to': '1234'},
+            'method': 'urn:oasis:names:tc:SAML:2.0:cm:bearer'}
+    }