More CLI tests & clearer bytes stuff

Ensuring that bytes are written correctly on all supported Python
versions, including when writing to stdout.

5 files changed, 178 insertions, 12 deletions

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 @@ from io import StringIO, BytesIO
 
 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 @@ def cli_args(*new_argv):
         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 @@ def cleanup_files(*filenames):
 
 
 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 @@ class KeygenTest(AbstractCliTest):
             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 @@ class KeygenTest(AbstractCliTest):
         # 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 @@ class TestByte(unittest.TestCase):
     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 @@ class DerTest(unittest.TestCase):
         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 @@ class DerTest(unittest.TestCase):
         key = rsa.key.PublicKey(3727264081, 65537)
         der = key.save_pkcs1('DER')
 
+        self.assertIsInstance(der, bytes)
         self.assertEqual(PUBLIC_DER, der)
 
 
@@ -125,6 +127,7 @@ class PemTest(unittest.TestCase):
         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 @@ class PemTest(unittest.TestCase):
         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 @@ class TestByteOutput(unittest.TestCase):
         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')))