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
|
"""
Unit tests for CLI entry points.
"""
import unittest
import sys
import functools
from contextlib import contextmanager
import os
from io import StringIO, BytesIO
import rsa
import rsa.cli
if sys.version_info[0] < 3:
IOClass = BytesIO
else:
IOClass = StringIO
@contextmanager
def captured_output():
"""Captures output to stdout and stderr"""
new_out, new_err = IOClass(), IOClass()
old_out, old_err = sys.stdout, sys.stderr
try:
sys.stdout, sys.stderr = new_out, new_err
yield sys.stdout, sys.stderr
finally:
sys.stdout, sys.stderr = old_out, old_err
@contextmanager
def cli_args(*new_argv):
"""Updates sys.argv[1:] for a single test."""
old_args = sys.argv[:]
sys.argv[1:] = [str(arg) for arg in new_argv]
try:
yield
finally:
sys.argv[1:] = old_args
def cleanup_files(*filenames):
"""Makes sure the files don't exist when the test runs, and deletes them afterward."""
def remove():
for fname in filenames:
if os.path.exists(fname):
os.unlink(fname)
def decorator(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
remove()
try:
return func(*args, **kwargs)
finally:
remove()
return wrapper
return decorator
class AbstractCliTest(unittest.TestCase):
def assertExits(self, status_code, func, *args, **kwargs):
try:
func(*args, **kwargs)
except SystemExit as ex:
if status_code == ex.code:
return
self.fail('SystemExit() raised by %r, but exited with code %i, expected %i' % (
func, ex.code, status_code))
else:
self.fail('SystemExit() not raised by %r' % func)
class KeygenTest(AbstractCliTest):
def test_keygen_no_args(self):
with cli_args():
self.assertExits(1, rsa.cli.keygen)
def test_keygen_priv_stdout(self):
with captured_output() as (out, err):
with cli_args(128):
rsa.cli.keygen()
lines = out.getvalue().splitlines()
self.assertEqual('-----BEGIN RSA PRIVATE KEY-----', lines[0])
self.assertEqual('-----END RSA PRIVATE KEY-----', lines[-1])
# The key size should be shown on stderr
self.assertTrue('128-bit key' in err.getvalue())
@cleanup_files('test_cli_privkey_out.pem')
def test_keygen_priv_out_pem(self):
with captured_output() as (out, err):
with cli_args('--out=test_cli_privkey_out.pem', '--form=PEM', 128):
rsa.cli.keygen()
# The key size should be shown on stderr
self.assertTrue('128-bit key' in err.getvalue())
# The output file should be shown on stderr
self.assertTrue('test_cli_privkey_out.pem' in err.getvalue())
# If we can load the file as PEM, it's good enough.
with open('test_cli_privkey_out.pem', 'rb') as pemfile:
rsa.PrivateKey.load_pkcs1(pemfile.read())
@cleanup_files('test_cli_privkey_out.der')
def test_keygen_priv_out_der(self):
with captured_output() as (out, err):
with cli_args('--out=test_cli_privkey_out.der', '--form=DER', 128):
rsa.cli.keygen()
# The key size should be shown on stderr
self.assertTrue('128-bit key' in err.getvalue())
# The output file should be shown on stderr
self.assertTrue('test_cli_privkey_out.der' in err.getvalue())
# If we can load the file as der, it's good enough.
with open('test_cli_privkey_out.der', 'rb') as derfile:
rsa.PrivateKey.load_pkcs1(derfile.read(), format='DER')
@cleanup_files('test_cli_privkey_out.pem', 'test_cli_pubkey_out.pem')
def test_keygen_pub_out_pem(self):
with captured_output() as (out, err):
with cli_args('--out=test_cli_privkey_out.pem',
'--pubout=test_cli_pubkey_out.pem',
'--form=PEM', 256):
rsa.cli.keygen()
# The key size should be shown on stderr
self.assertTrue('256-bit key' in err.getvalue())
# The output files should be shown on stderr
self.assertTrue('test_cli_privkey_out.pem' in err.getvalue())
self.assertTrue('test_cli_pubkey_out.pem' in err.getvalue())
# If we can load the file as PEM, it's good enough.
with open('test_cli_pubkey_out.pem', 'rb') as pemfile:
rsa.PublicKey.load_pkcs1(pemfile.read())
|
