#!/usr/bin/env python

"""
    chacha8_file_crypt_demo.py
     
    A Python demo program to encrypt/decrypt a file using ChaCha8.
    
    How I handle passwords and IVs is knowingly weak.  Remember, 
    this is only a demo.
     
    Larry Bugbee
    rev June 2010

"""

from chacha8 import ChaCha8
import sys, time, os, getpass
    
_version = '1'

#---------------------------------------------------------------

def usage():
    print '\n  usage:\n    python %s <filename> [<key>] \n' % sys.argv[0]

#---------------------------------------------------------------
# Simple file encryption/decryption

def crypt_file():

    suffix = '.cc8'
    chunksize = 64000

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    def writeallowed(filename):
        "return true if allowed to overwrite file, false otherwise"
        if os.path.exists(filename):
            prompt = '  "%s" exists.  overwrite? (y/[n]): ' % filename
            response = raw_input(prompt)
            if not (response and response.lower()[0] == 'y'):
                print '  original file preserved'
                return 0
        return 1
        
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    def savetofile(filename, content):
        "write content to file, unless disallowed"
        if writeallowed(filename):
            f = open(filename, 'wb')
            f.write(content)
            f.close()
        
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    def loadfmfile(filename):
        "load content from a file"
        f = open(filename, 'rb')
        content = f.read()
        f.close()
        return content
    
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    def bytestring(hex):
        "remove whitespace and convert hex string to a byte string"
        return hex.replace(' ', '').replace('\n', '').decode('hex')
    
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    def key1632(key):
        "pad or truncate keys to 16 or 32 bytes"
        newkey = key + chr(0)*32
        if len(key) <= 16:
            return newkey[:16]
        return newkey[:32]
    
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    def iv8(iv):
        "pad or truncate IV to 8 bytes"
        newiv = iv+chr(0)*8
        return (iv+chr(0)*8)[:8]
    
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    def makeNounce():
        """ create an 8 character unique number based on the lower
            bits of the system clock.  Add a .1 sec delay to help 
            ensure uniqueness.
            
            This is strictly a Q&D routine.  Serious programmers 
            will want something better.  
        """
        # sufficient delay on many processors to get a unique nonce
        time.sleep(0.1)
        nonce = ('%f' % (time.time())).replace('.', '')[-8:]
    
        '''
        # here are a few more options to further reduce probability 
        # of time as a nounce being reused.  For my purposes, time 
        # seems sufficient.
        # time, pid, sid, uid, and if OSX, Security Session Id
        # also consider urandom(), clock()
        ssid = os.environ.get('SECURITYSESSIONID', 'na')
        nonce += ' %d %d %d %s' % (os.getpid(), 
                                    os.getsid(os.getpid()), 
                                    os.getuid(),
                                    ssid,
                                   )
        '''
        return nonce
    
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    # analyze file name suffix
    filename = sys.argv[1]
    encrypt = not filename.endswith(suffix)
    if encrypt:                 # determined by file suffix    
        newfilename = filename+suffix
    else:
        newfilename = filename[:-len(suffix)]
    
    # pswd on commandline or do we need to enter without echo?
    if len(sys.argv) == 3:
        pswd = sys.argv[2]
    else:
        pswd = getpass.getpass('  enter your key: ')
        if encrypt:
            pswd2 = getpass.getpass('           again: ')
            if pswd != pswd2:
                print '  key entries not the same, try again'
                removeLibFile()
                sys.exit()
    key = key1632(pswd)         # pad or truncate as appropriate
    
    if writeallowed(newfilename):
        fi = open(filename, 'rb')
        fo = open(newfilename, 'wb')
        if encrypt:                 # determined by file suffix
            IV = iv8(makeNounce())  # create a IV, albeit funky
            fo.write(IV)            # put IV at front of file
            ctx = ChaCha8(key, IV)
            t0 = time.time()
            while 1:
                chunk = fi.read(chunksize)
                if not chunk: break
                ciphertext = ctx.encrypt(chunk)
                fo.write(ciphertext)
#            print '  %.3f seconds' % (time.time() - t0)
        else:
            IV = fi.read(8)         # get IV from front of file
            ctx = ChaCha8(key, IV)
            t0 = time.time()
            while 1:
                chunk = fi.read(chunksize)
                if not chunk: break
                plaintext = ctx.decrypt(chunk)
                fo.write(plaintext)
#            print '  %.3f seconds' % (time.time() - t0)
        fi.close()
        fo.close()
    
#---------------------------------------------------------------

if __name__ == '__main__':

    if len(sys.argv) < 2:
        usage()
        sys.exit()

    crypt_file()

#---------------------------------------------------------------
#---------------------------------------------------------------
#---------------------------------------------------------------