codekingpro/portable-devtools

1#  Copyright 2011 Sybren A. Stüvel <sybren@stuvel.eu>
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(
38        usage="usage: %prog [options] keysize",
39        description='Generates a new RSA key pair of "keysize" bits.',
40    )
41
42    parser.add_option(
43        "--pubout",
44        type="string",
45        help="Output filename for the public key. The public key is "
46        "not saved if this option is not present. You can use "
47        "pyrsa-priv2pub to create the public key file later.",
48    )
49
50    parser.add_option(
51        "-o",
52        "--out",
53        type="string",
54        help="Output filename for the private key. The key is "
55        "written to stdout if this option is not present.",
56    )
57
58    parser.add_option(
59        "--form",
60        help="key format of the private and public keys - default PEM",
61        choices=("PEM", "DER"),
62        default="PEM",
63    )
64
65    (cli, cli_args) = parser.parse_args(sys.argv[1:])
66
67    if len(cli_args) != 1:
68        parser.print_help()
69        raise SystemExit(1)
70
71    try:
72        keysize = int(cli_args[0])
73    except ValueError as ex:
74        parser.print_help()
75        print("Not a valid number: %s" % cli_args[0], file=sys.stderr)
76        raise SystemExit(1) from ex
77
78    print("Generating %i-bit key" % keysize, file=sys.stderr)
79    (pub_key, priv_key) = rsa.newkeys(keysize)
80
81    # Save public key
82    if cli.pubout:
83        print("Writing public key to %s" % cli.pubout, file=sys.stderr)
84        data = pub_key.save_pkcs1(format=cli.form)
85        with open(cli.pubout, "wb") as outfile:
86            outfile.write(data)
87
88    # Save private key
89    data = priv_key.save_pkcs1(format=cli.form)
90
91    if cli.out:
92        print("Writing private key to %s" % cli.out, file=sys.stderr)
93        with open(cli.out, "wb") as outfile:
94            outfile.write(data)
95    else:
96        print("Writing private key to stdout", file=sys.stderr)
97        sys.stdout.buffer.write(data)
98
99
100class CryptoOperation(metaclass=abc.ABCMeta):
101    """CLI callable that operates with input, output, and a key."""
102
103    keyname = "public"  # or 'private'
104    usage = "usage: %%prog [options] %(keyname)s_key"
105    description = ""
106    operation = "decrypt"
107    operation_past = "decrypted"
108    operation_progressive = "decrypting"
109    input_help = "Name of the file to %(operation)s. Reads from stdin if " "not specified."
110    output_help = (
111        "Name of the file to write the %(operation_past)s file "
112        "to. Written to stdout if this option is not present."
113    )
114    expected_cli_args = 1
115    has_output = True
116
117    key_class = rsa.PublicKey  # type: typing.Type[rsa.key.AbstractKey]
118
119    def __init__(self) -> None:
120        self.usage = self.usage % self.__class__.__dict__
121        self.input_help = self.input_help % self.__class__.__dict__
122        self.output_help = self.output_help % self.__class__.__dict__
123
124    @abc.abstractmethod
125    def perform_operation(
126        self, indata: bytes, key: rsa.key.AbstractKey, cli_args: Indexable
127    ) -> typing.Any:
128        """Performs the program's operation.
129
130        Implement in a subclass.
131
132        :returns: the data to write to the output.
133        """
134
135    def __call__(self) -> None:
136        """Runs the program."""
137
138        (cli, cli_args) = self.parse_cli()
139
140        key = self.read_key(cli_args[0], cli.keyform)
141
142        indata = self.read_infile(cli.input)
143
144        print(self.operation_progressive.title(), file=sys.stderr)
145        outdata = self.perform_operation(indata, key, cli_args)
146
147        if self.has_output:
148            self.write_outfile(outdata, cli.output)
149
150    def parse_cli(self) -> typing.Tuple[optparse.Values, typing.List[str]]:
151        """Parse the CLI options
152
153        :returns: (cli_opts, cli_args)
154        """
155
156        parser = optparse.OptionParser(usage=self.usage, description=self.description)
157
158        parser.add_option("-i", "--input", type="string", help=self.input_help)
159
160        if self.has_output:
161            parser.add_option("-o", "--output", type="string", help=self.output_help)
162
163        parser.add_option(
164            "--keyform",
165            help="Key format of the %s key - default PEM" % self.keyname,
166            choices=("PEM", "DER"),
167            default="PEM",
168        )
169
170        (cli, cli_args) = parser.parse_args(sys.argv[1:])
171
172        if len(cli_args) != self.expected_cli_args:
173            parser.print_help()
174            raise SystemExit(1)
175
176        return cli, cli_args
177
178    def read_key(self, filename: str, keyform: str) -> rsa.key.AbstractKey:
179        """Reads a public or private key."""
180
181        print("Reading %s key from %s" % (self.keyname, filename), file=sys.stderr)
182        with open(filename, "rb") as keyfile:
183            keydata = keyfile.read()
184
185        return self.key_class.load_pkcs1(keydata, keyform)
186
187    def read_infile(self, inname: str) -> bytes:
188        """Read the input file"""
189
190        if inname:
191            print("Reading input from %s" % inname, file=sys.stderr)
192            with open(inname, "rb") as infile:
193                return infile.read()
194
195        print("Reading input from stdin", file=sys.stderr)
196        return sys.stdin.buffer.read()
197
198    def write_outfile(self, outdata: bytes, outname: str) -> None:
199        """Write the output file"""
200
201        if outname:
202            print("Writing output to %s" % outname, file=sys.stderr)
203            with open(outname, "wb") as outfile:
204                outfile.write(outdata)
205        else:
206            print("Writing output to stdout", file=sys.stderr)
207            sys.stdout.buffer.write(outdata)
208
209
210class EncryptOperation(CryptoOperation):
211    """Encrypts a file."""
212
213    keyname = "public"
214    description = (
215        "Encrypts a file. The file must be shorter than the key " "length in order to be encrypted."
216    )
217    operation = "encrypt"
218    operation_past = "encrypted"
219    operation_progressive = "encrypting"
220
221    def perform_operation(
222        self, indata: bytes, pub_key: rsa.key.AbstractKey, cli_args: Indexable = ()
223    ) -> bytes:
224        """Encrypts files."""
225        assert isinstance(pub_key, rsa.key.PublicKey)
226        return rsa.encrypt(indata, pub_key)
227
228
229class DecryptOperation(CryptoOperation):
230    """Decrypts a file."""
231
232    keyname = "private"
233    description = (
234        "Decrypts a file. The original file must be shorter than "
235        "the key length in order to have been encrypted."
236    )
237    operation = "decrypt"
238    operation_past = "decrypted"
239    operation_progressive = "decrypting"
240    key_class = rsa.PrivateKey
241
242    def perform_operation(
243        self, indata: bytes, priv_key: rsa.key.AbstractKey, cli_args: Indexable = ()
244    ) -> bytes:
245        """Decrypts files."""
246        assert isinstance(priv_key, rsa.key.PrivateKey)
247        return rsa.decrypt(indata, priv_key)
248
249
250class SignOperation(CryptoOperation):
251    """Signs a file."""
252
253    keyname = "private"
254    usage = "usage: %%prog [options] private_key hash_method"
255    description = (
256        "Signs a file, outputs the signature. Choose the hash "
257        "method from %s" % ", ".join(HASH_METHODS)
258    )
259    operation = "sign"
260    operation_past = "signature"
261    operation_progressive = "Signing"
262    key_class = rsa.PrivateKey
263    expected_cli_args = 2
264
265    output_help = (
266        "Name of the file to write the signature to. Written "
267        "to stdout if this option is not present."
268    )
269
270    def perform_operation(
271        self, indata: bytes, priv_key: rsa.key.AbstractKey, cli_args: Indexable
272    ) -> bytes:
273        """Signs files."""
274        assert isinstance(priv_key, rsa.key.PrivateKey)
275
276        hash_method = cli_args[1]
277        if hash_method not in HASH_METHODS:
278            raise SystemExit("Invalid hash method, choose one of %s" % ", ".join(HASH_METHODS))
279
280        return rsa.sign(indata, priv_key, hash_method)
281
282
283class VerifyOperation(CryptoOperation):
284    """Verify a signature."""
285
286    keyname = "public"
287    usage = "usage: %%prog [options] public_key signature_file"
288    description = (
289        "Verifies a signature, exits with status 0 upon success, "
290        "prints an error message and exits with status 1 upon error."
291    )
292    operation = "verify"
293    operation_past = "verified"
294    operation_progressive = "Verifying"
295    key_class = rsa.PublicKey
296    expected_cli_args = 2
297    has_output = False
298
299    def perform_operation(
300        self, indata: bytes, pub_key: rsa.key.AbstractKey, cli_args: Indexable
301    ) -> None:
302        """Verifies files."""
303        assert isinstance(pub_key, rsa.key.PublicKey)
304
305        signature_file = cli_args[1]
306
307        with open(signature_file, "rb") as sigfile:
308            signature = sigfile.read()
309
310        try:
311            rsa.verify(indata, signature, pub_key)
312        except rsa.VerificationError as ex:
313            raise SystemExit("Verification failed.") from ex
314
315        print("Verification OK", file=sys.stderr)
316
317
318encrypt = EncryptOperation()
319decrypt = DecryptOperation()
320sign = SignOperation()
321verify = VerifyOperation()
322