1import os.path
2from openid.dh import DiffieHellman, strxor
3
4
5def test_strxor():
6    NUL = b'\x00'
7
8    cases = [
9        (NUL, NUL, NUL),
10        (b'\x01', NUL, b'\x01'),
11        (b'a', b'a', NUL),
12        (b'a', NUL, b'a'),
13        (b'abc', NUL * 3, b'abc'),
14        (b'x' * 10, NUL * 10, b'x' * 10),
15        (b'\x01', b'\x02', b'\x03'),
16        (b'\xf0', b'\x0f', b'\xff'),
17        (b'\xff', b'\x0f', b'\xf0'),
18    ]
19
20    for aa, bb, expected in cases:
21        actual = strxor(aa, bb)
22        assert actual == expected, (aa, bb, expected, actual)
23
24    exc_cases = [
25        ('', 'a'),
26        ('foo', 'ba'),
27        (NUL * 3, NUL * 4),
28        (''.join(map(chr, range(256))), ''.join(map(chr, range(128)))),
29    ]
30
31    for aa, bb in exc_cases:
32        try:
33            unexpected = strxor(aa, bb)
34        except ValueError:
35            pass
36        else:
37            assert False, 'Expected ValueError, got %r' % (unexpected, )
38
39
40def test1():
41    dh1 = DiffieHellman.fromDefaults()
42    dh2 = DiffieHellman.fromDefaults()
43    secret1 = dh1.getSharedSecret(dh2.public)
44    secret2 = dh2.getSharedSecret(dh1.public)
45    assert secret1 == secret2
46    return secret1
47
48
49def test_exchange():
50    s1 = test1()
51    s2 = test1()
52    assert s1 != s2
53
54
55def test_public():
56    f = open(os.path.join(os.path.dirname(__file__), 'dhpriv'))
57    dh = DiffieHellman.fromDefaults()
58    try:
59        for line in f:
60            parts = line.strip().split(' ')
61            dh._setPrivate(int(parts[0]))
62
63            assert dh.public == int(parts[1])
64    finally:
65        f.close()
66
67
68def test():
69    test_exchange()
70    test_public()
71    test_strxor()
72
73
74if __name__ == '__main__':
75    test()
76