1# -*- coding: utf-8 -*-
2#
3#  Copyright 2011 Sybren A. Stüvel <sybren@stuvel.eu>
4#
5#  Licensed under the Apache License, Version 2.0 (the "License");
6#  you may not use this file except in compliance with the License.
7#  You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11#  Unless required by applicable law or agreed to in writing, software
12#  distributed under the License is distributed on an "AS IS" BASIS,
13#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14#  See the License for the specific language governing permissions and
15#  limitations under the License.
16
17'''Commandline scripts.
18
19These scripts are called by the executables defined in setup.py.
20'''
21
22from __future__ import with_statement, print_function
23
24import abc
25import sys
26from optparse import OptionParser
27
28import rsa
29import rsa.bigfile
30import rsa.pkcs1
31
32HASH_METHODS = sorted(rsa.pkcs1.HASH_METHODS.keys())
33
34def keygen():
35    '''Key generator.'''
36
37    # Parse the CLI options
38    parser = OptionParser(usage='usage: %prog [options] keysize',
39            description='Generates a new RSA keypair of "keysize" bits.')
40
41    parser.add_option('--pubout', type='string',
42            help='Output filename for the public key. The public key is '
43            'not saved if this option is not present. You can use '
44            'pyrsa-priv2pub to create the public key file later.')
45
46    parser.add_option('-o', '--out', type='string',
47            help='Output filename for the private key. The key is '
48            'written to stdout if this option is not present.')
49
50    parser.add_option('--form',
51            help='key format of the private and public keys - default PEM',
52            choices=('PEM', 'DER'), default='PEM')
53
54    (cli, cli_args) = parser.parse_args(sys.argv[1:])
55
56    if len(cli_args) != 1:
57        parser.print_help()
58        raise SystemExit(1)
59
60    try:
61        keysize = int(cli_args[0])
62    except ValueError:
63        parser.print_help()
64        print('Not a valid number: %s' % cli_args[0], file=sys.stderr)
65        raise SystemExit(1)
66
67    print('Generating %i-bit key' % keysize, file=sys.stderr)
68    (pub_key, priv_key) = rsa.newkeys(keysize)
69
70
71    # Save public key
72    if cli.pubout:
73        print('Writing public key to %s' % cli.pubout, file=sys.stderr)
74        data = pub_key.save_pkcs1(format=cli.form)
75        with open(cli.pubout, 'wb') as outfile:
76            outfile.write(data)
77
78    # Save private key
79    data = priv_key.save_pkcs1(format=cli.form)
80
81    if cli.out:
82        print('Writing private key to %s' % cli.out, file=sys.stderr)
83        with open(cli.out, 'wb') as outfile:
84            outfile.write(data)
85    else:
86        print('Writing private key to stdout', file=sys.stderr)
87        sys.stdout.write(data)
88
89
90class CryptoOperation(object):
91    '''CLI callable that operates with input, output, and a key.'''
92
93    __metaclass__ = abc.ABCMeta
94
95    keyname = 'public' # or 'private'
96    usage = 'usage: %%prog [options] %(keyname)s_key'
97    description = None
98    operation = 'decrypt'
99    operation_past = 'decrypted'
100    operation_progressive = 'decrypting'
101    input_help = 'Name of the file to %(operation)s. Reads from stdin if ' \
102            'not specified.'
103    output_help = 'Name of the file to write the %(operation_past)s file ' \
104            'to. Written to stdout if this option is not present.'
105    expected_cli_args = 1
106    has_output = True
107
108    key_class = rsa.PublicKey
109
110    def __init__(self):
111        self.usage = self.usage % self.__class__.__dict__
112        self.input_help = self.input_help % self.__class__.__dict__
113        self.output_help = self.output_help % self.__class__.__dict__
114
115    @abc.abstractmethod
116    def perform_operation(self, indata, key, cli_args=None):
117        '''Performs the program's operation.
118
119        Implement in a subclass.
120
121        :returns: the data to write to the output.
122        '''
123
124    def __call__(self):
125        '''Runs the program.'''
126
127        (cli, cli_args) = self.parse_cli()
128
129        key = self.read_key(cli_args[0], cli.keyform)
130
131        indata = self.read_infile(cli.input)
132
133        print(self.operation_progressive.title(), file=sys.stderr)
134        outdata = self.perform_operation(indata, key, cli_args)
135
136        if self.has_output:
137            self.write_outfile(outdata, cli.output)
138
139    def parse_cli(self):
140        '''Parse the CLI options
141
142        :returns: (cli_opts, cli_args)
143        '''
144
145        parser = OptionParser(usage=self.usage, description=self.description)
146
147        parser.add_option('-i', '--input', type='string', help=self.input_help)
148
149        if self.has_output:
150            parser.add_option('-o', '--output', type='string', help=self.output_help)
151
152        parser.add_option('--keyform',
153                help='Key format of the %s key - default PEM' % self.keyname,
154                choices=('PEM', 'DER'), default='PEM')
155
156        (cli, cli_args) = parser.parse_args(sys.argv[1:])
157
158        if len(cli_args) != self.expected_cli_args:
159            parser.print_help()
160            raise SystemExit(1)
161
162        return (cli, cli_args)
163
164    def read_key(self, filename, keyform):
165        '''Reads a public or private key.'''
166
167        print('Reading %s key from %s' % (self.keyname, filename), file=sys.stderr)
168        with open(filename, 'rb') as keyfile:
169            keydata = keyfile.read()
170
171        return self.key_class.load_pkcs1(keydata, keyform)
172
173    def read_infile(self, inname):
174        '''Read the input file'''
175
176        if inname:
177            print('Reading input from %s' % inname, file=sys.stderr)
178            with open(inname, 'rb') as infile:
179                return infile.read()
180
181        print('Reading input from stdin', file=sys.stderr)
182        return sys.stdin.read()
183
184    def write_outfile(self, outdata, outname):
185        '''Write the output file'''
186
187        if outname:
188            print('Writing output to %s' % outname, file=sys.stderr)
189            with open(outname, 'wb') as outfile:
190                outfile.write(outdata)
191        else:
192            print('Writing output to stdout', file=sys.stderr)
193            sys.stdout.write(outdata)
194
195class EncryptOperation(CryptoOperation):
196    '''Encrypts a file.'''
197
198    keyname = 'public'
199    description = ('Encrypts a file. The file must be shorter than the key '
200            'length in order to be encrypted. For larger files, use the '
201            'pyrsa-encrypt-bigfile command.')
202    operation = 'encrypt'
203    operation_past = 'encrypted'
204    operation_progressive = 'encrypting'
205
206
207    def perform_operation(self, indata, pub_key, cli_args=None):
208        '''Encrypts files.'''
209
210        return rsa.encrypt(indata, pub_key)
211
212class DecryptOperation(CryptoOperation):
213    '''Decrypts a file.'''
214
215    keyname = 'private'
216    description = ('Decrypts a file. The original file must be shorter than '
217            'the key length in order to have been encrypted. For larger '
218            'files, use the pyrsa-decrypt-bigfile command.')
219    operation = 'decrypt'
220    operation_past = 'decrypted'
221    operation_progressive = 'decrypting'
222    key_class = rsa.PrivateKey
223
224    def perform_operation(self, indata, priv_key, cli_args=None):
225        '''Decrypts files.'''
226
227        return rsa.decrypt(indata, priv_key)
228
229class SignOperation(CryptoOperation):
230    '''Signs a file.'''
231
232    keyname = 'private'
233    usage = 'usage: %%prog [options] private_key hash_method'
234    description = ('Signs a file, outputs the signature. Choose the hash '
235            'method from %s' % ', '.join(HASH_METHODS))
236    operation = 'sign'
237    operation_past = 'signature'
238    operation_progressive = 'Signing'
239    key_class = rsa.PrivateKey
240    expected_cli_args = 2
241
242    output_help = ('Name of the file to write the signature to. Written '
243            'to stdout if this option is not present.')
244
245    def perform_operation(self, indata, priv_key, cli_args):
246        '''Decrypts files.'''
247
248        hash_method = cli_args[1]
249        if hash_method not in HASH_METHODS:
250            raise SystemExit('Invalid hash method, choose one of %s' %
251                    ', '.join(HASH_METHODS))
252
253        return rsa.sign(indata, priv_key, hash_method)
254
255class VerifyOperation(CryptoOperation):
256    '''Verify a signature.'''
257
258    keyname = 'public'
259    usage = 'usage: %%prog [options] public_key signature_file'
260    description = ('Verifies a signature, exits with status 0 upon success, '
261        'prints an error message and exits with status 1 upon error.')
262    operation = 'verify'
263    operation_past = 'verified'
264    operation_progressive = 'Verifying'
265    key_class = rsa.PublicKey
266    expected_cli_args = 2
267    has_output = False
268
269    def perform_operation(self, indata, pub_key, cli_args):
270        '''Decrypts files.'''
271
272        signature_file = cli_args[1]
273
274        with open(signature_file, 'rb') as sigfile:
275            signature = sigfile.read()
276
277        try:
278            rsa.verify(indata, signature, pub_key)
279        except rsa.VerificationError:
280            raise SystemExit('Verification failed.')
281
282        print('Verification OK', file=sys.stderr)
283
284
285class BigfileOperation(CryptoOperation):
286    '''CryptoOperation that doesn't read the entire file into memory.'''
287
288    def __init__(self):
289        CryptoOperation.__init__(self)
290
291        self.file_objects = []
292
293    def __del__(self):
294        '''Closes any open file handles.'''
295
296        for fobj in self.file_objects:
297            fobj.close()
298
299    def __call__(self):
300        '''Runs the program.'''
301
302        (cli, cli_args) = self.parse_cli()
303
304        key = self.read_key(cli_args[0], cli.keyform)
305
306        # Get the file handles
307        infile = self.get_infile(cli.input)
308        outfile = self.get_outfile(cli.output)
309
310        # Call the operation
311        print(self.operation_progressive.title(), file=sys.stderr)
312        self.perform_operation(infile, outfile, key, cli_args)
313
314    def get_infile(self, inname):
315        '''Returns the input file object'''
316
317        if inname:
318            print('Reading input from %s' % inname, file=sys.stderr)
319            fobj = open(inname, 'rb')
320            self.file_objects.append(fobj)
321        else:
322            print('Reading input from stdin', file=sys.stderr)
323            fobj = sys.stdin
324
325        return fobj
326
327    def get_outfile(self, outname):
328        '''Returns the output file object'''
329
330        if outname:
331            print('Will write output to %s' % outname, file=sys.stderr)
332            fobj = open(outname, 'wb')
333            self.file_objects.append(fobj)
334        else:
335            print('Will write output to stdout', file=sys.stderr)
336            fobj = sys.stdout
337
338        return fobj
339
340class EncryptBigfileOperation(BigfileOperation):
341    '''Encrypts a file to VARBLOCK format.'''
342
343    keyname = 'public'
344    description = ('Encrypts a file to an encrypted VARBLOCK file. The file '
345            'can be larger than the key length, but the output file is only '
346            'compatible with Python-RSA.')
347    operation = 'encrypt'
348    operation_past = 'encrypted'
349    operation_progressive = 'encrypting'
350
351    def perform_operation(self, infile, outfile, pub_key, cli_args=None):
352        '''Encrypts files to VARBLOCK.'''
353
354        return rsa.bigfile.encrypt_bigfile(infile, outfile, pub_key)
355
356class DecryptBigfileOperation(BigfileOperation):
357    '''Decrypts a file in VARBLOCK format.'''
358
359    keyname = 'private'
360    description = ('Decrypts an encrypted VARBLOCK file that was encrypted '
361            'with pyrsa-encrypt-bigfile')
362    operation = 'decrypt'
363    operation_past = 'decrypted'
364    operation_progressive = 'decrypting'
365    key_class = rsa.PrivateKey
366
367    def perform_operation(self, infile, outfile, priv_key, cli_args=None):
368        '''Decrypts a VARBLOCK file.'''
369
370        return rsa.bigfile.decrypt_bigfile(infile, outfile, priv_key)
371
372
373encrypt = EncryptOperation()
374decrypt = DecryptOperation()
375sign = SignOperation()
376verify = VerifyOperation()
377encrypt_bigfile = EncryptBigfileOperation()
378decrypt_bigfile = DecryptBigfileOperation()
379
380