aboutsummaryrefslogtreecommitdiff
path: root/aes_keywrap.py
blob: 382c3107ff8c9966b6cb41cfd28cb14c0f03a33e (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
#!/usr/bin/env python

"""
Python implementation of RFC 5649  AES Key Wrap With Padding,
using PyCrypto to supply the AES code.
"""

# Terminology mostly follows the RFC, including variable names.
#
# Block sizes get confusing: AES Key Wrap uses 64-bit blocks, not to
# be confused with AES, which uses 128-bit blocks.  In practice, this
# is less confusing than when reading the description, because we
# concatenate two 64-bit blocks just prior to performing an AES ECB
# operation, then immediately split the result back into a pair of
# 64-bit blocks.

class AESKeyWrapWithPadding(object):
    """
    Implementation of AES Key Wrap With Padding from RFC 5649.
    """

    class UnwrapError(Exception):
        "Something went wrong during unwrap."

    def __init__(self, key):
        from Crypto.Cipher import AES
        self.ctx = AES.new(key, AES.MODE_ECB)

    def _encrypt(self, b1, b2):
        aes_block = self.ctx.encrypt(b1 + b2)
        return aes_block[:8], aes_block[8:]

    def _decrypt(self, b1, b2):
        aes_block = self.ctx.decrypt(b1 + b2)
        return aes_block[:8], aes_block[8:]

    @staticmethod
    def _start_stop(start, stop):                    # Syntactic sugar
        step = -1 if start > stop else 1
        return xrange(start, stop + step, step)


    def wrap_key(self, Q):
        """
        Wrap a key according to RFC 5649 section 4.1.

        Q is the plaintext to be wrapped, a byte string.

        Returns C, the wrapped ciphertext.
        """

        from struct import pack, unpack

        m = len(Q)                              # Plaintext length
        if m % 8 != 0:                          # Pad Q if needed
            Q += "\x00" * (8 - (m % 8))
        R = [pack(">LL", 0xa65959a6, m)]        # Magic MSB(32,A), build LSB(32,A)
        R.extend(Q[i : i + 8]                   # Append Q
                 for i in xrange(0, len(Q), 8))

        n = len(R) - 1

        if n == 1:
            R[0], R[1] = self._encrypt(R[0], R[1])

        else:
            # RFC 3394 section 2.2.1
            for j in self._start_stop(0, 5):
                for i in self._start_stop(1, n):
                    R[0], R[i] = self._encrypt(R[0], R[i])
                    W0, W1 = unpack(">LL", R[0])
                    W1 ^= n * j + i
                    R[0] = pack(">LL", W0, W1)

        assert len(R) == (n + 1) and all(len(r) == 8 for r in R)
        return "".join(R)


    def unwrap_key(self, C):
        """
        Unwrap a key according to RFC 5649 section 4.2.

        C is the ciphertext to be unwrapped, a byte string

        Returns Q, the unwrapped plaintext.
        """

        from struct import pack, unpack

        if len(C) % 8 != 0:
            raise self.UnwrapError("Ciphertext length {} is not an integral number of blocks"
                                   .format(len(C)))

        n = (len(C) / 8) - 1
        R = [C[i : i + 8] for i in xrange(0, len(C), 8)]

        if n == 1:
            R[0], R[1] = self._decrypt(R[0], R[1])

        else:
            # RFC 3394 section 2.2.2 steps (1), (2), and part of (3)
            for j in self._start_stop(5, 0):
                for i in self._start_stop(n, 1):
                    W0, W1 = unpack(">LL", R[0])
                    W1 ^= n * j + i
                    R[0] = pack(">LL", W0, W1)
                    R[0], R[i] = self._decrypt(R[0], R[i])

        magic, m = unpack(">LL", R[0])

        if magic != 0xa65959a6:
            raise self.UnwrapError("Magic value in AIV should have been 0xa65959a6, was 0x{:02x}"
                              .format(magic))

        if m <= 8 * (n - 1) or m > 8 * n:
            raise self.UnwrapError("Length encoded in AIV out of range: m {}, n {}".format(m, n))

        R = "".join(R[1:])
        assert len(R) ==  8 * n

        if any(r != "\x00" for r in R[m:]):
            raise self.UnwrapError("Nonzero trailing bytes {}".format(R[m:].encode("hex")))

        return R[:m]


if __name__ == "__main__":

    # Test code from here down

    import unittest

    class TestAESKeyWrapWithPadding(unittest.TestCase):

        @staticmethod
        def bin2hex(bytes, sep = ":"):
            return sep.join("{:02x}".format(ord(b)) for b in bytes)

        @staticmethod
        def hex2bin(text):
            return text.translate(None, ": \t\n\r").decode("hex")

        def loopback_test(self, I):
            K = AESKeyWrapWithPadding(self.hex2bin("00:01:02:03:04:05:06:07:08:09:0a:0b:0c:0d:0e:0f"))
            C = K.wrap_key(I)
            O = K.unwrap_key(C)
            self.assertEqual(I, O, "Input and output plaintext did not match: {!r} <> {!r}".format(I, O))

        def rfc5649_test(self, K, Q, C):
            K = AESKeyWrapWithPadding(key = self.hex2bin(K))
            Q = self.hex2bin(Q)
            C = self.hex2bin(C)
            c = K.wrap_key(Q)
            q = K.unwrap_key(C)
            self.assertEqual(q, Q, "Input and output plaintext did not match: {} <> {}".format(self.bin2hex(Q), self.bin2hex(q)))
            self.assertEqual(c, C, "Input and output ciphertext did not match: {} <> {}".format(self.bin2hex(C), self.bin2hex(c)))

        def test_rfc5649_1(self):
            self.rfc5649_test(K = "5840df6e29b02af1 ab493b705bf16ea1 ae8338f4dcc176a8",
                              Q = "c37b7e6492584340 bed1220780894115 5068f738",
                              C = "138bdeaa9b8fa7fc 61f97742e72248ee 5ae6ae5360d1ae6a 5f54f373fa543b6a")

        def test_rfc5649_2(self):
            self.rfc5649_test(K = "5840df6e29b02af1 ab493b705bf16ea1 ae8338f4dcc176a8",
                              Q = "466f7250617369",
                              C = "afbeb0f07dfbf541 9200f2ccb50bb24f")

        def test_mangled_1(self):
            self.assertRaises(AESKeyWrapWithPadding.UnwrapError, self.rfc5649_test,
                              K = "5840df6e29b02af0 ab493b705bf16ea1 ae8338f4dcc176a8",
                              Q = "466f7250617368",
                              C = "afbeb0f07dfbf541 9200f2ccb50bb24f")

        def test_mangled_2(self):
            self.assertRaises(AESKeyWrapWithPadding.UnwrapError, self.rfc5649_test,
                              K = "5840df6e29b02af0 ab493b705bf16ea1 ae8338f4dcc176a8",
                              Q = "466f7250617368",
                              C = "afbeb0f07dfbf541 9200f2ccb50bb24f 0123456789abcdef")

        def test_mangled_3(self):
            self.assertRaises(AESKeyWrapWithPadding.UnwrapError, self.rfc5649_test,
                              K = "5840df6e29b02af1 ab493b705bf16ea1 ae8338f4dcc176a8",
                              Q = "c37b7e6492584340 bed1220780894115 5068f738",
                              C = "138bdeaa9b8fa7fc 61f97742e72248ee 5ae6ae5360d1ae6a")

        def test_loopback_1(self):
            self.loopback_test("!")

        def test_loopback_2(self):
            self.loopback_test("Yo!")

        def test_loopback_3(self):
            self.loopback_test("Hi, Mom")

        def test_loopback_4(self):
            self.loopback_test("1" * (64 / 8))

        def test_loopback_5(self):
            self.loopback_test("2" * (128 / 8))

        def test_loopback_6(self):
            self.loopback_test("3" * (256 / 8))

        def test_loopback_7(self):
            self.loopback_test("3.14159265358979323846264338327950288419716939937510")

        def test_loopback_8(self):
            self.loopback_test("3.14159265358979323846264338327950288419716939937510")

        def test_loopback_9(self):
            self.loopback_test("Hello!  My name is Inigo Montoya. You killed my AES key wrapper. Prepare to die.")

        def test_joachim_loopback(self):
            from os import urandom
            I = "31:32:33"
            K = AESKeyWrapWithPadding(urandom(256/8))
            C = K.wrap_key(I)
            O = K.unwrap_key(C)
            self.assertEqual(I, O, "Input and output plaintext did not match: {!r} <> {!r}".format(I, O))

    unittest.main(verbosity = 9)