More CLI tests & clearer bytes stuff
Ensuring that bytes are written correctly on all supported Python
versions, including when writing to stdout.
diff --git a/tests/__init__.py b/tests/__init__.py
index e69de29..85bafc7 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -0,0 +1,13 @@
+import unittest
+import sys
+
+
+if sys.hexversion < 0x2070000:
+ # Monkey-patch unittest.TestCase to add assertIsInstance on Python 2.6
+
+ def assertIsInstance(self, obj, cls, msg=None):
+ """Same as self.assertTrue(isinstance(obj, cls)), with a nicer default message."""
+ if not isinstance(obj, cls):
+ self.fail('%r is not an instance of %r but is a %r' % (obj, cls, type(obj)))
+
+ unittest.TestCase.assertIsInstance = assertIsInstance
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 0051aa1..511d25d 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -2,6 +2,8 @@
Unit tests for CLI entry points.
"""
+from __future__ import print_function
+
import unittest
import sys
import functools
@@ -12,22 +14,37 @@
import rsa
import rsa.cli
+from rsa._compat import b
if sys.version_info[0] < 3:
- IOClass = BytesIO
+ def make_buffer():
+ return BytesIO()
+
+
+ def get_bytes_out(out):
+ # Python 2.x writes 'str' to stdout:
+ return out.getvalue()
else:
- IOClass = StringIO
+ def make_buffer():
+ buf = StringIO()
+ buf.buffer = BytesIO()
+ return buf
+
+
+ def get_bytes_out(out):
+ # Python 3.x writes 'bytes' to stdout.buffer:
+ return out.buffer.getvalue()
@contextmanager
def captured_output():
"""Captures output to stdout and stderr"""
- new_out, new_err = IOClass(), IOClass()
+ new_out, new_err = make_buffer(), make_buffer()
old_out, old_err = sys.stdout, sys.stderr
try:
sys.stdout, sys.stderr = new_out, new_err
- yield sys.stdout, sys.stderr
+ yield new_out, new_err
finally:
sys.stdout, sys.stderr = old_out, old_err
@@ -45,13 +62,19 @@
sys.argv[1:] = old_args
+def remove_if_exists(fname):
+ """Removes a file if it exists."""
+
+ if os.path.exists(fname):
+ os.unlink(fname)
+
+
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)
+ remove_if_exists(fname)
def decorator(func):
@functools.wraps(func)
@@ -68,13 +91,33 @@
class AbstractCliTest(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ # Ensure there is a key to use
+ cls.pub_key, cls.priv_key = rsa.newkeys(512)
+ cls.pub_fname = '%s.pub' % cls.__name__
+ cls.priv_fname = '%s.key' % cls.__name__
+
+ with open(cls.pub_fname, 'wb') as outfile:
+ outfile.write(cls.pub_key.save_pkcs1())
+
+ with open(cls.priv_fname, 'wb') as outfile:
+ outfile.write(cls.priv_key.save_pkcs1())
+
+ @classmethod
+ def tearDownClass(cls):
+ if hasattr(cls, 'pub_fname'):
+ remove_if_exists(cls.pub_fname)
+ if hasattr(cls, 'priv_fname'):
+ remove_if_exists(cls.priv_fname)
+
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' % (
+ self.fail('SystemExit() raised by %r, but exited with code %r, expected %r' % (
func, ex.code, status_code))
else:
self.fail('SystemExit() not raised by %r' % func)
@@ -90,9 +133,9 @@
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])
+ lines = get_bytes_out(out).splitlines()
+ self.assertEqual(b('-----BEGIN RSA PRIVATE KEY-----'), lines[0])
+ self.assertEqual(b('-----END RSA PRIVATE KEY-----'), lines[-1])
# The key size should be shown on stderr
self.assertTrue('128-bit key' in err.getvalue())
@@ -147,3 +190,91 @@
# 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())
+
+
+class EncryptDecryptTest(AbstractCliTest):
+ def test_empty_decrypt(self):
+ with cli_args():
+ self.assertExits(1, rsa.cli.decrypt)
+
+ def test_empty_encrypt(self):
+ with cli_args():
+ self.assertExits(1, rsa.cli.encrypt)
+
+ @cleanup_files('encrypted.txt', 'cleartext.txt')
+ def test_encrypt_decrypt(self):
+ with open('cleartext.txt', 'wb') as outfile:
+ outfile.write(b'Hello cleartext RSA users!')
+
+ with cli_args('-i', 'cleartext.txt', '--out=encrypted.txt', self.pub_fname):
+ with captured_output():
+ rsa.cli.encrypt()
+
+ with cli_args('-i', 'encrypted.txt', self.priv_fname):
+ with captured_output() as (out, err):
+ rsa.cli.decrypt()
+
+ # We should have the original cleartext on stdout now.
+ output = get_bytes_out(out)
+ self.assertEqual(b('Hello cleartext RSA users!'), output)
+
+ @cleanup_files('encrypted.txt', 'cleartext.txt')
+ def test_encrypt_decrypt_unhappy(self):
+ with open('cleartext.txt', 'wb') as outfile:
+ outfile.write(b'Hello cleartext RSA users!')
+
+ with cli_args('-i', 'cleartext.txt', '--out=encrypted.txt', self.pub_fname):
+ with captured_output():
+ rsa.cli.encrypt()
+
+ # Change a few bytes in the encrypted stream.
+ with open('encrypted.txt', 'r+b') as encfile:
+ encfile.seek(40)
+ encfile.write(b'hahaha')
+
+ with cli_args('-i', 'encrypted.txt', self.priv_fname):
+ with captured_output() as (out, err):
+ self.assertRaises(rsa.DecryptionError, rsa.cli.decrypt)
+
+
+class SignVerifyTest(AbstractCliTest):
+ def test_empty_verify(self):
+ with cli_args():
+ self.assertExits(1, rsa.cli.verify)
+
+ def test_empty_sign(self):
+ with cli_args():
+ self.assertExits(1, rsa.cli.sign)
+
+ @cleanup_files('signature.txt', 'cleartext.txt')
+ def test_sign_verify(self):
+ with open('cleartext.txt', 'wb') as outfile:
+ outfile.write(b'Hello RSA users!')
+
+ with cli_args('-i', 'cleartext.txt', '--out=signature.txt', self.priv_fname, 'SHA-256'):
+ with captured_output():
+ rsa.cli.sign()
+
+ with cli_args('-i', 'cleartext.txt', self.pub_fname, 'signature.txt'):
+ with captured_output() as (out, err):
+ rsa.cli.verify()
+
+ self.assertFalse(b'Verification OK' in get_bytes_out(out))
+
+ @cleanup_files('signature.txt', 'cleartext.txt')
+ def test_sign_verify_unhappy(self):
+ with open('cleartext.txt', 'wb') as outfile:
+ outfile.write(b'Hello RSA users!')
+
+ with cli_args('-i', 'cleartext.txt', '--out=signature.txt', self.priv_fname, 'SHA-256'):
+ with captured_output():
+ rsa.cli.sign()
+
+ # Change a few bytes in the cleartext file.
+ with open('cleartext.txt', 'r+b') as encfile:
+ encfile.seek(6)
+ encfile.write(b'DSA')
+
+ with cli_args('-i', 'cleartext.txt', self.pub_fname, 'signature.txt'):
+ with captured_output() as (out, err):
+ self.assertExits('Verification failed.', rsa.cli.verify)
diff --git a/tests/test_compat.py b/tests/test_compat.py
index 8cbf101..74d6f53 100644
--- a/tests/test_compat.py
+++ b/tests/test_compat.py
@@ -16,8 +16,9 @@
import unittest
import struct
+import sys
-from rsa._compat import is_bytes, byte
+from rsa._compat import is_bytes, byte, b
class TestByte(unittest.TestCase):
@@ -30,3 +31,6 @@
def test_raises_StructError_on_overflow(self):
self.assertRaises(struct.error, byte, 256)
self.assertRaises(struct.error, byte, -1)
+
+ def test_byte_literal(self):
+ self.assertIsInstance(b('abc'), bytes)
diff --git a/tests/test_load_save_keys.py b/tests/test_load_save_keys.py
index 6f374cf..0caa067 100644
--- a/tests/test_load_save_keys.py
+++ b/tests/test_load_save_keys.py
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-'''Unittest for saving and loading keys.'''
+"""Unittest for saving and loading keys."""
import base64
import unittest
@@ -89,6 +89,7 @@
key = rsa.key.PrivateKey(3727264081, 65537, 3349121513, 65063, 57287)
der = key.save_pkcs1('DER')
+ self.assertIsInstance(der, bytes)
self.assertEqual(PRIVATE_DER, der)
def test_load_public_key(self):
@@ -105,6 +106,7 @@
key = rsa.key.PublicKey(3727264081, 65537)
der = key.save_pkcs1('DER')
+ self.assertIsInstance(der, bytes)
self.assertEqual(PUBLIC_DER, der)
@@ -125,6 +127,7 @@
key = rsa.key.PrivateKey(3727264081, 65537, 3349121513, 65063, 57287)
pem = key.save_pkcs1('PEM')
+ self.assertIsInstance(pem, bytes)
self.assertEqual(CLEAN_PRIVATE_PEM, pem)
def test_load_public_key(self):
@@ -141,6 +144,7 @@
key = rsa.key.PublicKey(3727264081, 65537)
pem = key.save_pkcs1('PEM')
+ self.assertIsInstance(pem, bytes)
self.assertEqual(CLEAN_PUBLIC_PEM, pem)
def test_load_from_disk(self):
diff --git a/tests/test_pem.py b/tests/test_pem.py
index 61a66fc..3e03ab0 100644
--- a/tests/test_pem.py
+++ b/tests/test_pem.py
@@ -86,3 +86,17 @@
key = rsa.key.PrivateKey.load_pkcs1(private_key_pem)
self.assertTrue(is_bytes(key.save_pkcs1(format='DER')))
self.assertTrue(is_bytes(key.save_pkcs1(format='PEM')))
+
+
+class TestByteInput(unittest.TestCase):
+ """Tests that PEM and DER can be loaded from bytes."""
+
+ def test_bytes_public(self):
+ key = rsa.key.PublicKey.load_pkcs1_openssl_pem(public_key_pem.encode('ascii'))
+ self.assertTrue(is_bytes(key.save_pkcs1(format='DER')))
+ self.assertTrue(is_bytes(key.save_pkcs1(format='PEM')))
+
+ def test_bytes_private(self):
+ key = rsa.key.PrivateKey.load_pkcs1(private_key_pem.encode('ascii'))
+ self.assertTrue(is_bytes(key.save_pkcs1(format='DER')))
+ self.assertTrue(is_bytes(key.save_pkcs1(format='PEM')))
