1#  Copyright 2011 Sybren A. Stüvel <[email protected]>
2#
3#  Licensed under the Apache License, Version 2.0 (the "License");
4#  you may not use this file except in compliance with the License.
5#  You may obtain a copy of the License at
6#
7#      https://www.apache.org/licenses/LICENSE-2.0
8#
9#  Unless required by applicable law or agreed to in writing, software
10#  distributed under the License is distributed on an "AS IS" BASIS,
11#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12#  See the License for the specific language governing permissions and
13#  limitations under the License.
14
15"""Commandline scripts.
16
17These scripts are called by the executables defined in setup.py.
18"""
19
20import abc
21import sys
22import typing
23import optparse
24
25import rsa
26import rsa.key
27import rsa.pkcs1
28
29HASH_METHODS = sorted(rsa.pkcs1.HASH_METHODS.keys())
30Indexable = typing.Union[typing.Tuple, typing.List[str]]
31
32
33def keygen() -> None:
34    """Key generator."""
35
36    # Parse the CLI options
37    parser = optparse.OptionParser(usage='usage: %prog [options] keysize',
38                                   description='Generates a new RSA keypair of "keysize" bits.')
39
40    parser.add_option('--pubout', type='string',
41                      help='Output filename for the public key. The public key is '
42                           'not saved if this option is not present. You can use '
43                           'pyrsa-priv2pub to create the public key file later.')
44
45    parser.add_option('-o', '--out', type='string',
46                      help='Output filename for the private key. The key is '
47                           'written to stdout if this option is not present.')
48
49    parser.add_option('--form',
50                      help='key format of the private and public keys - default PEM',
51                      choices=('PEM', 'DER'), default='PEM')
52
53    (cli, cli_args) = parser.parse_args(sys.argv[1:])
54
55    if len(cli_args) != 1:
56        parser.print_help()
57        raise SystemExit(1)
58
59    try:
60        keysize = int(cli_args[0])
61    except ValueError:
62        parser.print_help()
63        print('Not a valid number: %s' % cli_args[0], file=sys.stderr)
64        raise SystemExit(1)
65
66    print('Generating %i-bit key' % keysize, file=sys.stderr)
67    (pub_key, priv_key) = rsa.newkeys(keysize)
68
69    # Save public key
70    if cli.pubout:
71        print('Writing public key to %s' % cli.pubout, file=sys.stderr)
72        data = pub_key.save_pkcs1(format=cli.form)
73        with open(cli.pubout, 'wb') as outfile:
74            outfile.write(data)
75
76    # Save private key
77    data = priv_key.save_pkcs1(format=cli.form)
78
79    if cli.out:
80        print('Writing private key to %s' % cli.out, file=sys.stderr)
81        with open(cli.out, 'wb') as outfile:
82            outfile.write(data)
83    else:
84        print('Writing private key to stdout', file=sys.stderr)
85        sys.stdout.buffer.write(data)
86
87
88class CryptoOperation(metaclass=abc.ABCMeta):
89    """CLI callable that operates with input, output, and a key."""
90
91    keyname = 'public'  # or 'private'
92    usage = 'usage: %%prog [options] %(keyname)s_key'
93    description = ''
94    operation = 'decrypt'
95    operation_past = 'decrypted'
96    operation_progressive = 'decrypting'
97    input_help = 'Name of the file to %(operation)s. Reads from stdin if ' \
98                 'not specified.'
99    output_help = 'Name of the file to write the %(operation_past)s file ' \
100                  'to. Written to stdout if this option is not present.'
101    expected_cli_args = 1
102    has_output = True
103
104    key_class = rsa.PublicKey  # type: typing.Type[rsa.key.AbstractKey]
105
106    def __init__(self) -> None:
107        self.usage = self.usage % self.__class__.__dict__
108        self.input_help = self.input_help % self.__class__.__dict__
109        self.output_help = self.output_help % self.__class__.__dict__
110
111    @abc.abstractmethod
112    def perform_operation(self, indata: bytes, key: rsa.key.AbstractKey,
113                          cli_args: Indexable) -> typing.Any:
114        """Performs the program's operation.
115
116        Implement in a subclass.
117
118        :returns: the data to write to the output.
119        """
120
121    def __call__(self) -> None:
122        """Runs the program."""
123
124        (cli, cli_args) = self.parse_cli()
125
126        key = self.read_key(cli_args[0], cli.keyform)
127
128        indata = self.read_infile(cli.input)
129
130        print(self.operation_progressive.title(), file=sys.stderr)
131        outdata = self.perform_operation(indata, key, cli_args)
132
133        if self.has_output:
134            self.write_outfile(outdata, cli.output)
135
136    def parse_cli(self) -> typing.Tuple[optparse.Values, typing.List[str]]:
137        """Parse the CLI options
138
139        :returns: (cli_opts, cli_args)
140        """
141
142        parser = optparse.OptionParser(usage=self.usage, description=self.description)
143
144        parser.add_option('-i', '--input', type='string', help=self.input_help)
145
146        if self.has_output:
147            parser.add_option('-o', '--output', type='string', help=self.output_help)
148
149        parser.add_option('--keyform',
150                          help='Key format of the %s key - default PEM' % self.keyname,
151                          choices=('PEM', 'DER'), default='PEM')
152
153        (cli, cli_args) = parser.parse_args(sys.argv[1:])
154
155        if len(cli_args) != self.expected_cli_args:
156            parser.print_help()
157            raise SystemExit(1)
158
159        return cli, cli_args
160
161    def read_key(self, filename: str, keyform: str) -> rsa.key.AbstractKey:
162        """Reads a public or private key."""
163
164        print('Reading %s key from %s' % (self.keyname, filename), file=sys.stderr)
165        with open(filename, 'rb') as keyfile:
166            keydata = keyfile.read()
167
168        return self.key_class.load_pkcs1(keydata, keyform)
169
170    def read_infile(self, inname: str) -> bytes:
171        """Read the input file"""
172
173        if inname:
174            print('Reading input from %s' % inname, file=sys.stderr)
175            with open(inname, 'rb') as infile:
176                return infile.read()
177
178        print('Reading input from stdin', file=sys.stderr)
179        return sys.stdin.buffer.read()
180
181    def write_outfile(self, outdata: bytes, outname: str) -> None:
182        """Write the output file"""
183
184        if outname:
185            print('Writing output to %s' % outname, file=sys.stderr)
186            with open(outname, 'wb') as outfile:
187                outfile.write(outdata)
188        else:
189            print('Writing output to stdout', file=sys.stderr)
190            sys.stdout.buffer.write(outdata)
191
192
193class EncryptOperation(CryptoOperation):
194    """Encrypts a file."""
195
196    keyname = 'public'
197    description = ('Encrypts a file. The file must be shorter than the key '
198                   'length in order to be encrypted.')
199    operation = 'encrypt'
200    operation_past = 'encrypted'
201    operation_progressive = 'encrypting'
202
203    def perform_operation(self, indata: bytes, pub_key: rsa.key.AbstractKey,
204                          cli_args: Indexable = ()) -> bytes:
205        """Encrypts files."""
206        assert isinstance(pub_key, rsa.key.PublicKey)
207        return rsa.encrypt(indata, pub_key)
208
209
210class DecryptOperation(CryptoOperation):
211    """Decrypts a file."""
212
213    keyname = 'private'
214    description = ('Decrypts a file. The original file must be shorter than '
215                   'the key length in order to have been encrypted.')
216    operation = 'decrypt'
217    operation_past = 'decrypted'
218    operation_progressive = 'decrypting'
219    key_class = rsa.PrivateKey
220
221    def perform_operation(self, indata: bytes, priv_key: rsa.key.AbstractKey,
222                          cli_args: Indexable = ()) -> bytes:
223        """Decrypts files."""
224        assert isinstance(priv_key, rsa.key.PrivateKey)
225        return rsa.decrypt(indata, priv_key)
226
227
228class SignOperation(CryptoOperation):
229    """Signs a file."""
230
231    keyname = 'private'
232    usage = 'usage: %%prog [options] private_key hash_method'
233    description = ('Signs a file, outputs the signature. Choose the hash '
234                   'method from %s' % ', '.join(HASH_METHODS))
235    operation = 'sign'
236    operation_past = 'signature'
237    operation_progressive = 'Signing'
238    key_class = rsa.PrivateKey
239    expected_cli_args = 2
240
241    output_help = ('Name of the file to write the signature to. Written '
242                   'to stdout if this option is not present.')
243
244    def perform_operation(self, indata: bytes, priv_key: rsa.key.AbstractKey,
245                          cli_args: Indexable) -> bytes:
246        """Signs files."""
247        assert isinstance(priv_key, rsa.key.PrivateKey)
248
249        hash_method = cli_args[1]
250        if hash_method not in HASH_METHODS:
251            raise SystemExit('Invalid hash method, choose one of %s' %
252                             ', '.join(HASH_METHODS))
253
254        return rsa.sign(indata, priv_key, hash_method)
255
256
257class VerifyOperation(CryptoOperation):
258    """Verify a signature."""
259
260    keyname = 'public'
261    usage = 'usage: %%prog [options] public_key signature_file'
262    description = ('Verifies a signature, exits with status 0 upon success, '
263                   'prints an error message and exits with status 1 upon error.')
264    operation = 'verify'
265    operation_past = 'verified'
266    operation_progressive = 'Verifying'
267    key_class = rsa.PublicKey
268    expected_cli_args = 2
269    has_output = False
270
271    def perform_operation(self, indata: bytes, pub_key: rsa.key.AbstractKey,
272                          cli_args: Indexable) -> None:
273        """Verifies files."""
274        assert isinstance(pub_key, rsa.key.PublicKey)
275
276        signature_file = cli_args[1]
277
278        with open(signature_file, 'rb') as sigfile:
279            signature = sigfile.read()
280
281        try:
282            rsa.verify(indata, signature, pub_key)
283        except rsa.VerificationError:
284            raise SystemExit('Verification failed.')
285
286        print('Verification OK', file=sys.stderr)
287
288
289encrypt = EncryptOperation()
290decrypt = DecryptOperation()
291sign = SignOperation()
292verify = VerifyOperation()
293