sign_file.py (2744B)
1 #! /usr/bin/env python 2 3 """ 4 A script that creates signed Python files. 5 """ 6 7 from __future__ import print_function 8 9 import os 10 import argparse 11 import subprocess 12 13 14 def get_file_encoding(filename): 15 """ 16 Get the file encoding for the file with the given filename 17 """ 18 19 with open(filename, 'rb') as fp: 20 # The encoding is usually specified on the second line 21 txt = fp.read().splitlines()[1] 22 txt = txt.decode('utf-8') 23 if 'encoding' in txt: 24 encoding = txt.split()[-1] 25 else: 26 encoding = 'utf-8' # default 27 28 return str(encoding) 29 30 31 def sign_file_and_get_sig(filename, encoding): 32 """ 33 Sign the file and get the signature 34 """ 35 36 cmd = 'gpg -bass {}'.format(filename) 37 ret = subprocess.Popen(cmd, shell=True).wait() 38 print ('-> %r' % cmd) 39 40 if ret: 41 raise ValueError('Could not sign the file!') 42 43 with open('{}.asc'.format(filename), 'rb') as fp: 44 sig = fp.read() 45 46 try: 47 os.remove('{}.asc'.format(filename)) 48 except OSError: 49 pass 50 51 sig = sig.decode(encoding) 52 sig = sig.replace('\r', '').replace('\n', '\\n') 53 sig = sig.encode(encoding) 54 55 return sig 56 57 58 def sign_original_file(filename, encoding): 59 """ 60 Sign the original file 61 """ 62 63 sig = sign_file_and_get_sig(filename, encoding) 64 65 with open(filename, 'ab') as outfile: 66 outfile.write('#'.encode(encoding)) 67 outfile.write(sig) 68 outfile.write('\n'.encode(encoding)) 69 70 71 def create_signed_file(filename, encoding): 72 """ 73 Create a signed file 74 """ 75 76 sig = sign_file_and_get_sig(filename, encoding) 77 name, extension = os.path.splitext(filename) 78 79 new_file_name = '{}_signed{}'.format(name, extension) 80 81 with open(new_file_name, 'wb') as outfile, \ 82 open(filename, 'rb') as infile: 83 txt = infile.read() 84 outfile.write(txt) 85 outfile.write('#'.encode(encoding)) 86 outfile.write(sig) 87 outfile.write('\n'.encode(encoding)) 88 89 90 def parse_args(): 91 parser = argparse.ArgumentParser() 92 93 parser.add_argument('filenames', action='store', nargs='+', 94 help='Files you wish to sign') 95 96 parser.add_argument('--overwrite', action='store_true', 97 dest='overwrite', default=False, 98 help='Overwrite the original file' 99 ' (sign the original file)') 100 101 opts = parser.parse_args() 102 103 return opts 104 105 106 if __name__ == '__main__': 107 opts = parse_args() 108 109 for filename in opts.filenames: 110 encoding = get_file_encoding(filename) 111 if opts.overwrite: 112 sign_original_file(filename, encoding) 113 else: 114 create_signed_file(filename, encoding) 115