Added encrypt and decrypt CLI commands

2 files changed, 150 insertions, 1 deletions

diff --git a/rsa/cli.py b/rsa/cli.py
index 764360e..e291f99 100644
--- a/rsa/cli.py
+++ b/rsa/cli.py
@@ -19,6 +19,7 @@
 These scripts are called by the executables defined in setup.py.
 '''
 
+import abc
 import sys
 from optparse import OptionParser
 
@@ -80,4 +81,150 @@ def keygen():
         print >>sys.stderr, 'Writing private key to stdout'
         sys.stdout.write(data)
 
-# vim: foldnestmax=1
+
+class CryptoOperation(object):
+    '''CLI callable that operates with input, output, and a key.'''
+
+    __metaclass__ = abc.ABCMeta
+
+    keyname = 'public' # or 'private'
+    usage = 'usage: %%prog [options] %(keyname)s_key'
+    description = None
+    operation = 'decrypt'
+    operation_past = 'decrypted'
+    operation_progressive = 'decrypting'
+    input_help = 'Name of the file to %(operation)s. Reads from stdin if ' \
+            '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.'
+
+    key_class = rsa.PublicKey
+
+    def __init__(self):
+        self.usage = self.usage % self.__class__.__dict__
+        self.input_help = self.input_help % self.__class__.__dict__
+        self.output_help = self.output_help % self.__class__.__dict__
+
+    @abc.abstractmethod
+    def perform_operation(self, indata, key):
+        '''Performs the program's operation.
+
+        Implement in a subclass.
+
+        :returns: the data to write to the output.
+        '''
+
+    def __call__(self):
+        '''Runs the program.'''
+
+        (cli, cli_args) = self.parse_cli()
+
+        key = self.read_key(cli_args[0], cli.keyform)
+
+        indata = self.read_infile(cli.input)
+
+        print >>sys.stderr, self.operation_progressive.title()
+        outdata = self.perform_operation(indata, key)
+
+        self.write_outfile(outdata, cli.output)
+
+    def parse_cli(self):
+        '''Parse the CLI options
+        
+        :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.add_option('--input', type='string',
+                help='Name of the file to encrypt. Reads from stdin if '
+                     'not specified.')
+
+        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.')
+
+        parser.add_option('--keyform',
+                help='Key format of the key - default PEM',
+                choices=('PEM', 'DER'), default='PEM')
+
+        (cli, cli_args) = parser.parse_args(sys.argv[1:])
+
+        if len(cli_args) != 1:
+            parser.print_help()
+            raise SystemExit(1)
+
+        return (cli, cli_args)
+
+    def read_key(self, filename, keyform):
+        '''Reads a public or private key.'''
+
+        print >>sys.stderr, 'Reading %s key from %s' % (self.keyname, filename)
+        with open(filename) as keyfile:
+            keydata = keyfile.read()
+
+        return self.key_class.load_pkcs1(keydata, keyform)
+    
+    def read_infile(self, inname):
+        '''Read the input file'''
+
+        if inname:
+            print >>sys.stderr, 'Reading input from %s' % inname
+            with open(inname, 'rb') as infile:
+                return infile.read()
+
+        print >>sys.stderr, 'Reading input from stdin'
+        return sys.stdin.read()
+
+    def write_outfile(self, outdata, outname):
+        '''Write the output file'''
+
+        if outname:
+            print >>sys.stderr, 'Writing output to %s' % outname
+            with open(outname, 'wb') as outfile:
+                outfile.write(outdata)
+        else:
+            print >>sys.stderr, 'Writing output to stdout'
+            sys.stdout.write(outdata)
+
+class EncryptOperation(CryptoOperation):
+    '''Encrypts a file.'''
+
+    keyname = 'public'
+    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.')
+    operation = 'encrypt'
+    operation_past = 'encrypted'
+    operation_progressive = 'encrypting'
+
+
+    def perform_operation(self, indata, pub_key):
+        '''Encrypts files.'''
+
+        return rsa.encrypt(indata, pub_key)
+
+class DecryptOperation(CryptoOperation):
+    '''Decrypts a file.'''
+
+    keyname = 'private'
+    description = ('Decrypts a file. The original file must be shorter than '
+            'the key length in order to have been encrypted. For larger '
+            'files, use the pyrsa-decrypt-bigfile command.')
+    operation = 'decrypt'
+    operation_past = 'decrypted'
+    operation_progressive = 'decrypting'
+    key_class = rsa.PrivateKey
+
+    def perform_operation(self, indata, priv_key):
+        '''Decrypts files.'''
+
+        return rsa.decrypt(indata, priv_key)
+
+encrypt = EncryptOperation()
+decrypt = DecryptOperation()
+
+
diff --git a/setup.py b/setup.py
index 8f3a726..436374a 100755
--- a/setup.py
+++ b/setup.py
@@ -35,6 +35,8 @@ setup(name='rsa',
     entry_points={ 'console_scripts': [
         'pyrsa-priv2pub = rsa.util:private_to_public',
         'pyrsa-keygen = rsa.cli:keygen',
+        'pyrsa-encrypt = rsa.cli:encrypt',
+        'pyrsa-decrypt = rsa.cli:decrypt',
     ]},
 
 )