Added verify and sign commands, fixed some bugs in other commands

3 files changed, 80 insertions, 20 deletions

diff --git a/rsa/__init__.py b/rsa/__init__.py
index a4147b2..1e60233 100644
--- a/rsa/__init__.py
+++ b/rsa/__init__.py
@@ -32,7 +32,8 @@ __date__ = "2010-02-08"
 __version__ = '3.0-beta0'
 
 from rsa.key import newkeys, PrivateKey, PublicKey
-from rsa.pkcs1 import encrypt, decrypt, sign, verify
+from rsa.pkcs1 import encrypt, decrypt, sign, verify, DecryptionError, \
+    VerificationError
 
 # Do doctest if we're run directly
 if __name__ == "__main__":
@@ -40,5 +41,5 @@ if __name__ == "__main__":
     doctest.testmod()
 
 __all__ = ["newkeys", "encrypt", "decrypt", "sign", "verify", 'PublicKey',
-    'PrivateKey']
+    'PrivateKey', 'DecryptionError', 'VerificationError']
 
diff --git a/rsa/cli.py b/rsa/cli.py
index e291f99..4ddaf35 100644
--- a/rsa/cli.py
+++ b/rsa/cli.py
@@ -25,6 +25,9 @@ from optparse import OptionParser
 
 import rsa
 import rsa.bigfile
+import rsa.pkcs1
+
+HASH_METHODS = sorted(rsa.pkcs1.HASH_METHODS.keys())
 
 def keygen():
     '''Key generator.'''
@@ -97,6 +100,8 @@ class CryptoOperation(object):
             'not specified.'
     output_help = 'Name of the file to write the %(operation_past)s file ' \
             'to. Written to stdout if this option is not present.'
+    expected_cli_args = 1
+    has_output = True
 
     key_class = rsa.PublicKey
 
@@ -106,7 +111,7 @@ class CryptoOperation(object):
         self.output_help = self.output_help % self.__class__.__dict__
 
     @abc.abstractmethod
-    def perform_operation(self, indata, key):
+    def perform_operation(self, indata, key, cli_args=None):
         '''Performs the program's operation.
 
         Implement in a subclass.
@@ -124,9 +129,10 @@ class CryptoOperation(object):
         indata = self.read_infile(cli.input)
 
         print >>sys.stderr, self.operation_progressive.title()
-        outdata = self.perform_operation(indata, key)
+        outdata = self.perform_operation(indata, key, cli_args)
 
-        self.write_outfile(outdata, cli.output)
+        if self.has_output:
+            self.write_outfile(outdata, cli.output)
 
     def parse_cli(self):
         '''Parse the CLI options
@@ -134,26 +140,20 @@ class CryptoOperation(object):
         :returns: (cli_opts, cli_args)
         '''
 
-        parser = OptionParser(usage='usage: %prog [options] public_key',
-                description='Encrypts a file. The file must be shorter than the '
-                'key length in order to be encrypted. For larger files, use the '
-                'pyrsa-encrypt-bigfile command.')
+        parser = OptionParser(usage=self.usage, description=self.description)
         
-        parser.add_option('--input', type='string',
-                help='Name of the file to encrypt. Reads from stdin if '
-                     'not specified.')
+        parser.add_option('--input', type='string', help=self.input_help)
 
-        parser.add_option('--output', type='string',
-                help='Name of the file to write the encrypted file to. '
-                     'Written to stdout if this option is not present.')
+        if self.has_output:
+            parser.add_option('--output', type='string', help=self.output_help)
 
         parser.add_option('--keyform',
-                help='Key format of the key - default PEM',
+                help='Key format of the %s key - default PEM' % self.keyname,
                 choices=('PEM', 'DER'), default='PEM')
 
         (cli, cli_args) = parser.parse_args(sys.argv[1:])
 
-        if len(cli_args) != 1:
+        if len(cli_args) != self.expected_cli_args:
             parser.print_help()
             raise SystemExit(1)
 
@@ -202,7 +202,7 @@ class EncryptOperation(CryptoOperation):
     operation_progressive = 'encrypting'
 
 
-    def perform_operation(self, indata, pub_key):
+    def perform_operation(self, indata, pub_key, cli_args=None):
         '''Encrypts files.'''
 
         return rsa.encrypt(indata, pub_key)
@@ -219,12 +219,69 @@ class DecryptOperation(CryptoOperation):
     operation_progressive = 'decrypting'
     key_class = rsa.PrivateKey
 
-    def perform_operation(self, indata, priv_key):
+    def perform_operation(self, indata, priv_key, cli_args=None):
         '''Decrypts files.'''
 
         return rsa.decrypt(indata, priv_key)
 
+class SignOperation(CryptoOperation):
+    '''Signs a file.'''
+
+    keyname = 'private'
+    usage = 'usage: %%prog [options] private_key hash_method'
+    description = ('Signs a file, outputs the signature. Choose the hash '
+            'method from %s' % ', '.join(HASH_METHODS))
+    operation = 'sign'
+    operation_past = 'signature'
+    operation_progressive = 'Signing'
+    key_class = rsa.PrivateKey
+    expected_cli_args = 2
+
+    output_help = ('Name of the file to write the signature to. Written '
+            'to stdout if this option is not present.')
+
+    def perform_operation(self, indata, priv_key, cli_args):
+        '''Decrypts files.'''
+
+        hash_method = cli_args[1]
+        if hash_method not in HASH_METHODS:
+            raise SystemExit('Invalid hash method, choose one of %s' % 
+                    ', '.join(HASH_METHODS))
+
+        return rsa.sign(indata, priv_key, hash_method)
+
+class VerifyOperation(CryptoOperation):
+    '''Verify a signature.'''
+
+    keyname = 'public'
+    usage = 'usage: %%prog [options] private_key signature_file'
+    description = ('Verifies a signature, exits with status 0 upon success, '
+        'prints an error message and exits with status 1 upon error.')
+    operation = 'verify'
+    operation_past = 'verified'
+    operation_progressive = 'Verifying'
+    key_class = rsa.PublicKey
+    expected_cli_args = 2
+    has_output = False
+
+    def perform_operation(self, indata, pub_key, cli_args):
+        '''Decrypts files.'''
+
+        signature_file = cli_args[1]
+        
+        with open(signature_file, 'rb') as sigfile:
+            signature = sigfile.read()
+
+        try:
+            rsa.verify(indata, signature, pub_key)
+        except rsa.VerificationError:
+            raise SystemExit('Verification failed.')
+
+        print >>sys.stderr, 'Verification OK'
+
+
 encrypt = EncryptOperation()
 decrypt = DecryptOperation()
-
+sign = SignOperation()
+verify = VerifyOperation()
 
diff --git a/setup.py b/setup.py
index 436374a..6286e58 100755
--- a/setup.py
+++ b/setup.py
@@ -37,6 +37,8 @@ setup(name='rsa',
         'pyrsa-keygen = rsa.cli:keygen',
         'pyrsa-encrypt = rsa.cli:encrypt',
         'pyrsa-decrypt = rsa.cli:decrypt',
+        'pyrsa-sign = rsa.cli:sign',
+        'pyrsa-verify = rsa.cli:verify',
     ]},
 
 )