cli.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383
  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.bigfile
  25. import rsa.pkcs1
  26. HASH_METHODS = sorted(rsa.pkcs1.HASH_METHODS.keys())
  27. def keygen():
  28. """Key generator."""
  29. # Parse the CLI options
  30. parser = OptionParser(usage='usage: %prog [options] keysize',
  31. description='Generates a new RSA keypair of "keysize" bits.')
  32. parser.add_option('--pubout', type='string',
  33. help='Output filename for the public key. The public key is '
  34. 'not saved if this option is not present. You can use '
  35. 'pyrsa-priv2pub to create the public key file later.')
  36. parser.add_option('-o', '--out', type='string',
  37. help='Output filename for the private key. The key is '
  38. 'written to stdout if this option is not present.')
  39. parser.add_option('--form',
  40. help='key format of the private and public keys - default PEM',
  41. choices=('PEM', 'DER'), default='PEM')
  42. (cli, cli_args) = parser.parse_args(sys.argv[1:])
  43. if len(cli_args) != 1:
  44. parser.print_help()
  45. raise SystemExit(1)
  46. try:
  47. keysize = int(cli_args[0])
  48. except ValueError:
  49. parser.print_help()
  50. print('Not a valid number: %s' % cli_args[0], file=sys.stderr)
  51. raise SystemExit(1)
  52. print('Generating %i-bit key' % keysize, file=sys.stderr)
  53. (pub_key, priv_key) = rsa.newkeys(keysize)
  54. # Save public key
  55. if cli.pubout:
  56. print('Writing public key to %s' % cli.pubout, file=sys.stderr)
  57. data = pub_key.save_pkcs1(format=cli.form)
  58. with open(cli.pubout, 'wb') as outfile:
  59. outfile.write(data)
  60. # Save private key
  61. data = priv_key.save_pkcs1(format=cli.form)
  62. if cli.out:
  63. print('Writing private key to %s' % cli.out, file=sys.stderr)
  64. with open(cli.out, 'wb') as outfile:
  65. outfile.write(data)
  66. else:
  67. print('Writing private key to stdout', file=sys.stderr)
  68. sys.stdout.write(data)
  69. class CryptoOperation(object):
  70. """CLI callable that operates with input, output, and a key."""
  71. __metaclass__ = abc.ABCMeta
  72. keyname = 'public' # or 'private'
  73. usage = 'usage: %%prog [options] %(keyname)s_key'
  74. description = None
  75. operation = 'decrypt'
  76. operation_past = 'decrypted'
  77. operation_progressive = 'decrypting'
  78. input_help = 'Name of the file to %(operation)s. Reads from stdin if ' \
  79. 'not specified.'
  80. output_help = 'Name of the file to write the %(operation_past)s file ' \
  81. 'to. Written to stdout if this option is not present.'
  82. expected_cli_args = 1
  83. has_output = True
  84. key_class = rsa.PublicKey
  85. def __init__(self):
  86. self.usage = self.usage % self.__class__.__dict__
  87. self.input_help = self.input_help % self.__class__.__dict__
  88. self.output_help = self.output_help % self.__class__.__dict__
  89. @abc.abstractmethod
  90. def perform_operation(self, indata, key, cli_args=None):
  91. """Performs the program's operation.
  92. Implement in a subclass.
  93. :returns: the data to write to the output.
  94. """
  95. def __call__(self):
  96. """Runs the program."""
  97. (cli, cli_args) = self.parse_cli()
  98. key = self.read_key(cli_args[0], cli.keyform)
  99. indata = self.read_infile(cli.input)
  100. print(self.operation_progressive.title(), file=sys.stderr)
  101. outdata = self.perform_operation(indata, key, cli_args)
  102. if self.has_output:
  103. self.write_outfile(outdata, cli.output)
  104. def parse_cli(self):
  105. """Parse the CLI options
  106. :returns: (cli_opts, cli_args)
  107. """
  108. parser = OptionParser(usage=self.usage, description=self.description)
  109. parser.add_option('-i', '--input', type='string', help=self.input_help)
  110. if self.has_output:
  111. parser.add_option('-o', '--output', type='string', help=self.output_help)
  112. parser.add_option('--keyform',
  113. help='Key format of the %s key - default PEM' % self.keyname,
  114. choices=('PEM', 'DER'), default='PEM')
  115. (cli, cli_args) = parser.parse_args(sys.argv[1:])
  116. if len(cli_args) != self.expected_cli_args:
  117. parser.print_help()
  118. raise SystemExit(1)
  119. return cli, cli_args
  120. def read_key(self, filename, keyform):
  121. """Reads a public or private key."""
  122. print('Reading %s key from %s' % (self.keyname, filename), file=sys.stderr)
  123. with open(filename, 'rb') as keyfile:
  124. keydata = keyfile.read()
  125. return self.key_class.load_pkcs1(keydata, keyform)
  126. def read_infile(self, inname):
  127. """Read the input file"""
  128. if inname:
  129. print('Reading input from %s' % inname, file=sys.stderr)
  130. with open(inname, 'rb') as infile:
  131. return infile.read()
  132. print('Reading input from stdin', file=sys.stderr)
  133. return sys.stdin.read()
  134. def write_outfile(self, outdata, outname):
  135. """Write the output file"""
  136. if outname:
  137. print('Writing output to %s' % outname, file=sys.stderr)
  138. with open(outname, 'wb') as outfile:
  139. outfile.write(outdata)
  140. else:
  141. print('Writing output to stdout', file=sys.stderr)
  142. sys.stdout.write(outdata)
  143. class EncryptOperation(CryptoOperation):
  144. """Encrypts a file."""
  145. keyname = 'public'
  146. description = ('Encrypts a file. The file must be shorter than the key '
  147. 'length in order to be encrypted. For larger files, use the '
  148. 'pyrsa-encrypt-bigfile command.')
  149. operation = 'encrypt'
  150. operation_past = 'encrypted'
  151. operation_progressive = 'encrypting'
  152. def perform_operation(self, indata, pub_key, cli_args=None):
  153. """Encrypts files."""
  154. return rsa.encrypt(indata, pub_key)
  155. class DecryptOperation(CryptoOperation):
  156. """Decrypts a file."""
  157. keyname = 'private'
  158. description = ('Decrypts a file. The original file must be shorter than '
  159. 'the key length in order to have been encrypted. For larger '
  160. 'files, use the pyrsa-decrypt-bigfile command.')
  161. operation = 'decrypt'
  162. operation_past = 'decrypted'
  163. operation_progressive = 'decrypting'
  164. key_class = rsa.PrivateKey
  165. def perform_operation(self, indata, priv_key, cli_args=None):
  166. """Decrypts files."""
  167. return rsa.decrypt(indata, priv_key)
  168. class SignOperation(CryptoOperation):
  169. """Signs a file."""
  170. keyname = 'private'
  171. usage = 'usage: %%prog [options] private_key hash_method'
  172. description = ('Signs a file, outputs the signature. Choose the hash '
  173. 'method from %s' % ', '.join(HASH_METHODS))
  174. operation = 'sign'
  175. operation_past = 'signature'
  176. operation_progressive = 'Signing'
  177. key_class = rsa.PrivateKey
  178. expected_cli_args = 2
  179. output_help = ('Name of the file to write the signature to. Written '
  180. 'to stdout if this option is not present.')
  181. def perform_operation(self, indata, priv_key, cli_args):
  182. """Signs files."""
  183. hash_method = cli_args[1]
  184. if hash_method not in HASH_METHODS:
  185. raise SystemExit('Invalid hash method, choose one of %s' %
  186. ', '.join(HASH_METHODS))
  187. return rsa.sign(indata, priv_key, hash_method)
  188. class VerifyOperation(CryptoOperation):
  189. """Verify a signature."""
  190. keyname = 'public'
  191. usage = 'usage: %%prog [options] public_key signature_file'
  192. description = ('Verifies a signature, exits with status 0 upon success, '
  193. 'prints an error message and exits with status 1 upon error.')
  194. operation = 'verify'
  195. operation_past = 'verified'
  196. operation_progressive = 'Verifying'
  197. key_class = rsa.PublicKey
  198. expected_cli_args = 2
  199. has_output = False
  200. def perform_operation(self, indata, pub_key, cli_args):
  201. """Verifies files."""
  202. signature_file = cli_args[1]
  203. with open(signature_file, 'rb') as sigfile:
  204. signature = sigfile.read()
  205. try:
  206. rsa.verify(indata, signature, pub_key)
  207. except rsa.VerificationError:
  208. raise SystemExit('Verification failed.')
  209. print('Verification OK', file=sys.stderr)
  210. class BigfileOperation(CryptoOperation):
  211. """CryptoOperation that doesn't read the entire file into memory."""
  212. def __init__(self):
  213. CryptoOperation.__init__(self)
  214. self.file_objects = []
  215. def __del__(self):
  216. """Closes any open file handles."""
  217. for fobj in self.file_objects:
  218. fobj.close()
  219. def __call__(self):
  220. """Runs the program."""
  221. (cli, cli_args) = self.parse_cli()
  222. key = self.read_key(cli_args[0], cli.keyform)
  223. # Get the file handles
  224. infile = self.get_infile(cli.input)
  225. outfile = self.get_outfile(cli.output)
  226. # Call the operation
  227. print(self.operation_progressive.title(), file=sys.stderr)
  228. self.perform_operation(infile, outfile, key, cli_args)
  229. def get_infile(self, inname):
  230. """Returns the input file object"""
  231. if inname:
  232. print('Reading input from %s' % inname, file=sys.stderr)
  233. fobj = open(inname, 'rb')
  234. self.file_objects.append(fobj)
  235. else:
  236. print('Reading input from stdin', file=sys.stderr)
  237. fobj = sys.stdin
  238. return fobj
  239. def get_outfile(self, outname):
  240. """Returns the output file object"""
  241. if outname:
  242. print('Will write output to %s' % outname, file=sys.stderr)
  243. fobj = open(outname, 'wb')
  244. self.file_objects.append(fobj)
  245. else:
  246. print('Will write output to stdout', file=sys.stderr)
  247. fobj = sys.stdout
  248. return fobj
  249. class EncryptBigfileOperation(BigfileOperation):
  250. """Encrypts a file to VARBLOCK format."""
  251. keyname = 'public'
  252. description = ('Encrypts a file to an encrypted VARBLOCK file. The file '
  253. 'can be larger than the key length, but the output file is only '
  254. 'compatible with Python-RSA.')
  255. operation = 'encrypt'
  256. operation_past = 'encrypted'
  257. operation_progressive = 'encrypting'
  258. def perform_operation(self, infile, outfile, pub_key, cli_args=None):
  259. """Encrypts files to VARBLOCK."""
  260. return rsa.bigfile.encrypt_bigfile(infile, outfile, pub_key)
  261. class DecryptBigfileOperation(BigfileOperation):
  262. """Decrypts a file in VARBLOCK format."""
  263. keyname = 'private'
  264. description = ('Decrypts an encrypted VARBLOCK file that was encrypted '
  265. 'with pyrsa-encrypt-bigfile')
  266. operation = 'decrypt'
  267. operation_past = 'decrypted'
  268. operation_progressive = 'decrypting'
  269. key_class = rsa.PrivateKey
  270. def perform_operation(self, infile, outfile, priv_key, cli_args=None):
  271. """Decrypts a VARBLOCK file."""
  272. return rsa.bigfile.decrypt_bigfile(infile, outfile, priv_key)
  273. encrypt = EncryptOperation()
  274. decrypt = DecryptOperation()
  275. sign = SignOperation()
  276. verify = VerifyOperation()
  277. encrypt_bigfile = EncryptBigfileOperation()
  278. decrypt_bigfile = DecryptBigfileOperation()