cryptopals/pkcs1.py at master · AdrS/cryptopals · GitHub

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
import random

def num_bytes(n):
	'''returns number of bytes needed to represent number'''
	b = 0
	while n:
		b += 1
		n >>= 1
	if b & 7:
		return (b >> 3) + 1
	else:
		return b >> 3

def pkcs1_pad(m, mod_len, rand=random.SystemRandom()):
	'''pad m according to PKCS1 so that it's same length as modulus'''
	pad_len = mod_len - len(m) - 3
	assert(pad_len >= 8)

	#padding starts with 00 02
	parts = ['\x00\x02']

	#all remaining padding (except last byte) is random nonzero bytes
	parts.extend([chr(rand.randrange(1, 256)) for _ in range(pad_len)])

	#padding ends with 00
	parts.append('\x00')
	parts.append(m)

	padded = ''.join(parts)
	assert(len(padded) == mod_len)

	return padded

def pkcs1_unpad(m, mod_len):
	'''removes padding returns none if padding is invalid'''
	if len(m) != mod_len: return
	if m[0] != '\x00' or m[1] != '\x02': return

	#first nonzero byte (after 00 02) marks end of padding
	i = 2
	while i < len(m) and m[i] != '\x00':
		i += 1

	#check that 00 byte was actually found
	if i >= len(m): return

	pad_len = i - 2
	#must have at least 8 bytes of random padding
	if pad_len < 8: return

	return m[3 + pad_len:]

def test_num_bytes():
	assert(num_bytes(0) == 0)
	assert(num_bytes(1) == 1)
	assert(num_bytes(7) == 1)
	assert(num_bytes(255) == 1)
	assert(num_bytes(256) == 2)
	assert(num_bytes((1<<16) - 1) == 2)

def test_pkcs1_unpad():
	#wrong length
	assert(not pkcs1_unpad('\x00\x0212345678\x00m', 13))

	#wrong prefix
	assert(not pkcs1_unpad('\x02\x0212345678\x00m', 12))

	# < 8 bytes of random padding
	assert(not pkcs1_unpad('\x00\x021234567\x00mm', 12))

	#no end of padding byte
	assert(not pkcs1_unpad('\x00\x0212345678pp', 12))

	#empty message
	assert(pkcs1_unpad('\x00\x0212345678p\x00', 12) == '')

	# > 8 bytes of random padding
	assert(pkcs1_unpad('\x00\x0212345678ppppppp\x00mm', 20) == 'mm')

def sanity_test():
	for m in ('we', 'all', 'live', 'yellow', 'beatles'):
		for mod_len in [18, 32, 64, 128]:
			assert(pkcs1_unpad(pkcs1_pad(m, mod_len), mod_len) == m)

if __name__ == '__main__':
	test_num_bytes()
	test_pkcs1_unpad()
	sanity_test()