def setUp(self):
self.key = ElGamalKey()
self.key.generate_for_nist(384)
self.test_file = 'lipsum.txt'
self.encr_file = 'test.enc'
self.decr_file = 'test.dec'
class TestElGamal(unittest.TestCase):
def setUp(self):
self.key = ElGamalKey()
self.key.generate_for_nist(384)
self.test_file = 'lipsum.txt'
self.encr_file = 'test.enc'
self.decr_file = 'test.dec'
def test_encryption_decryption(self):
encrypt(self.key, self.test_file, self.encr_file)
decrypt(self.key, self.encr_file, self.decr_file)
initial = open(self.test_file).read()
decrypted = open(self.decr_file).read()
self.assertEqual(initial, decrypted)
# clean up
for f in ['test.enc', 'test.dec']:
if os.path.exists(f):
os.remove(f)
def main():
parser = get_arg_parser()
options, args = parser.parse_args()
key = ElGamalKey() # create an empty key
if options.keyname: # generate keypair
if options.encrypt is not None:
print 'Cannot generate key and encrypt/decrypt at the same time. Aborting...'
else:
try:
key.generate_for_nist(options.size)
key.to_file('%s.pub' % options.keyname, 'public')
key.to_file('%s.priv' % options.keyname, 'private')
except ValueError:
print 'Can only use NIST curves. See help for sizes. Aborting...'
elif options.encrypt is not None:
# helper (DRY)
def _helper(op_type):
action = {
'encrypt': (encrypt, 'public'),
'decrypt': (decrypt, 'private')
}[op_type]
# need a key
if not options.key:
print 'No %s key provided. Aborting...' % action[1]
else:
try:
key.from_file(options.key, action[1])
if options.filein and os.path.exists(options.filein): # need an input file
action[0](key, options.filein, options.fileout)
else:
print 'No input file or file does not exist. Aborting...'
except: # bad key
print 'Wrong key format. Expecting %s key. Abort...' % action[1]
if options.encrypt: # encrypt
_helper('encrypt')
else: # decrypt
_helper('decrypt')
else:
parser.print_help()