cli.py 9.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288
  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. # https://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. """Commandline scripts.
  17. These scripts are called by the executables defined in setup.py.
  18. """
  19. from __future__ import with_statement, print_function
  20. import abc
  21. import sys
  22. from optparse import OptionParser
  23. import rsa
  24. import rsa.pkcs1
  25. HASH_METHODS = sorted(rsa.pkcs1.HASH_METHODS.keys())
  26. def keygen():
  27. """Key generator."""
  28. # Parse the CLI options
  29. parser = OptionParser(usage='usage: %prog [options] keysize',
  30. description='Generates a new RSA keypair of "keysize" bits.')
  31. parser.add_option('--pubout', type='string',
  32. help='Output filename for the public key. The public key is '
  33. 'not saved if this option is not present. You can use '
  34. 'pyrsa-priv2pub to create the public key file later.')
  35. parser.add_option('-o', '--out', type='string',
  36. help='Output filename for the private key. The key is '
  37. 'written to stdout if this option is not present.')
  38. parser.add_option('--form',
  39. help='key format of the private and public keys - default PEM',
  40. choices=('PEM', 'DER'), default='PEM')
  41. (cli, cli_args) = parser.parse_args(sys.argv[1:])
  42. if len(cli_args) != 1:
  43. parser.print_help()
  44. raise SystemExit(1)
  45. try:
  46. keysize = int(cli_args[0])
  47. except ValueError:
  48. parser.print_help()
  49. print('Not a valid number: %s' % cli_args[0], file=sys.stderr)
  50. raise SystemExit(1)
  51. print('Generating %i-bit key' % keysize, file=sys.stderr)
  52. (pub_key, priv_key) = rsa.newkeys(keysize)
  53. # Save public key
  54. if cli.pubout:
  55. print('Writing public key to %s' % cli.pubout, file=sys.stderr)
  56. data = pub_key.save_pkcs1(format=cli.form)
  57. with open(cli.pubout, 'wb') as outfile:
  58. outfile.write(data)
  59. # Save private key
  60. data = priv_key.save_pkcs1(format=cli.form)
  61. if cli.out:
  62. print('Writing private key to %s' % cli.out, file=sys.stderr)
  63. with open(cli.out, 'wb') as outfile:
  64. outfile.write(data)
  65. else:
  66. print('Writing private key to stdout', file=sys.stderr)
  67. rsa._compat.write_to_stdout(data)
  68. class CryptoOperation(object):
  69. """CLI callable that operates with input, output, and a key."""
  70. __metaclass__ = abc.ABCMeta
  71. keyname = 'public' # or 'private'
  72. usage = 'usage: %%prog [options] %(keyname)s_key'
  73. description = None
  74. operation = 'decrypt'
  75. operation_past = 'decrypted'
  76. operation_progressive = 'decrypting'
  77. input_help = 'Name of the file to %(operation)s. Reads from stdin if ' \
  78. 'not specified.'
  79. output_help = 'Name of the file to write the %(operation_past)s file ' \
  80. 'to. Written to stdout if this option is not present.'
  81. expected_cli_args = 1
  82. has_output = True
  83. key_class = rsa.PublicKey
  84. def __init__(self):
  85. self.usage = self.usage % self.__class__.__dict__
  86. self.input_help = self.input_help % self.__class__.__dict__
  87. self.output_help = self.output_help % self.__class__.__dict__
  88. @abc.abstractmethod
  89. def perform_operation(self, indata, key, cli_args):
  90. """Performs the program's operation.
  91. Implement in a subclass.
  92. :returns: the data to write to the output.
  93. """
  94. def __call__(self):
  95. """Runs the program."""
  96. (cli, cli_args) = self.parse_cli()
  97. key = self.read_key(cli_args[0], cli.keyform)
  98. indata = self.read_infile(cli.input)
  99. print(self.operation_progressive.title(), file=sys.stderr)
  100. outdata = self.perform_operation(indata, key, cli_args)
  101. if self.has_output:
  102. self.write_outfile(outdata, cli.output)
  103. def parse_cli(self):
  104. """Parse the CLI options
  105. :returns: (cli_opts, cli_args)
  106. """
  107. parser = OptionParser(usage=self.usage, description=self.description)
  108. parser.add_option('-i', '--input', type='string', help=self.input_help)
  109. if self.has_output:
  110. parser.add_option('-o', '--output', type='string', help=self.output_help)
  111. parser.add_option('--keyform',
  112. help='Key format of the %s key - default PEM' % self.keyname,
  113. choices=('PEM', 'DER'), default='PEM')
  114. (cli, cli_args) = parser.parse_args(sys.argv[1:])
  115. if len(cli_args) != self.expected_cli_args:
  116. parser.print_help()
  117. raise SystemExit(1)
  118. return cli, cli_args
  119. def read_key(self, filename, keyform):
  120. """Reads a public or private key."""
  121. print('Reading %s key from %s' % (self.keyname, filename), file=sys.stderr)
  122. with open(filename, 'rb') as keyfile:
  123. keydata = keyfile.read()
  124. return self.key_class.load_pkcs1(keydata, keyform)
  125. def read_infile(self, inname):
  126. """Read the input file"""
  127. if inname:
  128. print('Reading input from %s' % inname, file=sys.stderr)
  129. with open(inname, 'rb') as infile:
  130. return infile.read()
  131. print('Reading input from stdin', file=sys.stderr)
  132. return sys.stdin.read()
  133. def write_outfile(self, outdata, outname):
  134. """Write the output file"""
  135. if outname:
  136. print('Writing output to %s' % outname, file=sys.stderr)
  137. with open(outname, 'wb') as outfile:
  138. outfile.write(outdata)
  139. else:
  140. print('Writing output to stdout', file=sys.stderr)
  141. rsa._compat.write_to_stdout(outdata)
  142. class EncryptOperation(CryptoOperation):
  143. """Encrypts a file."""
  144. keyname = 'public'
  145. description = ('Encrypts a file. The file must be shorter than the key '
  146. 'length in order to be encrypted.')
  147. operation = 'encrypt'
  148. operation_past = 'encrypted'
  149. operation_progressive = 'encrypting'
  150. def perform_operation(self, indata, pub_key, cli_args=None):
  151. """Encrypts files."""
  152. return rsa.encrypt(indata, pub_key)
  153. class DecryptOperation(CryptoOperation):
  154. """Decrypts a file."""
  155. keyname = 'private'
  156. description = ('Decrypts a file. The original file must be shorter than '
  157. 'the key length in order to have been encrypted.')
  158. operation = 'decrypt'
  159. operation_past = 'decrypted'
  160. operation_progressive = 'decrypting'
  161. key_class = rsa.PrivateKey
  162. def perform_operation(self, indata, priv_key, cli_args=None):
  163. """Decrypts files."""
  164. return rsa.decrypt(indata, priv_key)
  165. class SignOperation(CryptoOperation):
  166. """Signs a file."""
  167. keyname = 'private'
  168. usage = 'usage: %%prog [options] private_key hash_method'
  169. description = ('Signs a file, outputs the signature. Choose the hash '
  170. 'method from %s' % ', '.join(HASH_METHODS))
  171. operation = 'sign'
  172. operation_past = 'signature'
  173. operation_progressive = 'Signing'
  174. key_class = rsa.PrivateKey
  175. expected_cli_args = 2
  176. output_help = ('Name of the file to write the signature to. Written '
  177. 'to stdout if this option is not present.')
  178. def perform_operation(self, indata, priv_key, cli_args):
  179. """Signs files."""
  180. hash_method = cli_args[1]
  181. if hash_method not in HASH_METHODS:
  182. raise SystemExit('Invalid hash method, choose one of %s' %
  183. ', '.join(HASH_METHODS))
  184. return rsa.sign(indata, priv_key, hash_method)
  185. class VerifyOperation(CryptoOperation):
  186. """Verify a signature."""
  187. keyname = 'public'
  188. usage = 'usage: %%prog [options] public_key signature_file'
  189. description = ('Verifies a signature, exits with status 0 upon success, '
  190. 'prints an error message and exits with status 1 upon error.')
  191. operation = 'verify'
  192. operation_past = 'verified'
  193. operation_progressive = 'Verifying'
  194. key_class = rsa.PublicKey
  195. expected_cli_args = 2
  196. has_output = False
  197. def perform_operation(self, indata, pub_key, cli_args):
  198. """Verifies files."""
  199. signature_file = cli_args[1]
  200. with open(signature_file, 'rb') as sigfile:
  201. signature = sigfile.read()
  202. try:
  203. rsa.verify(indata, signature, pub_key)
  204. except rsa.VerificationError:
  205. raise SystemExit('Verification failed.')
  206. print('Verification OK', file=sys.stderr)
  207. encrypt = EncryptOperation()
  208. decrypt = DecryptOperation()
  209. sign = SignOperation()
  210. verify = VerifyOperation()