#!/usr/bin/env python
# coding: utf-8

"""
    hotSalsa.py
    
        [ This is an earlier version, intended for use 
          with pySalsa20-v3 or later.  This is NOT the 
          self-contained version of hotSalsa. ]
    
    A small program to quickly and easily encrypt small files 
    using the Salsa20 algorithm.  ........
      http://en.wikipedia.org/wiki/Salsa20
      http://en.wikipedia.org/wiki/Daniel_Bernstein
      http://cr.yp.to/djb.html
      http://www.ecrypt.eu.org/stream/salsa20p3.html
    
    A new encrypted file is named the same as the original 
    file but with a '.s20' suffix added.  A decrypted file 
    is written without the '.s20' suffix.  If the output file 
    already exists and is about to be overwritten, the user is
    asked for permission.
    
    hotSalsa does not chunk large files so file size is 
    limited by available memory.  Sizes in the many-megabyte 
    range should work fine.  A minor mod would be required to 
    process files of larger sizes.  
    
      Usage:
        python hotSalsa.py <filename> [<key>]
    
    The key may be any string of up to 32 characters.  Be sure 
    to enclose the key in quotes if it has imbedded spaces or 
    special characters.  If the key is not on the commandline, 
    the user will be asked to type it in (no echo).
    
    Salsa20 wants the key to be exactly 128 or 256 bits (16 or 
    32 bytes).  This program, hotSalsa, will pad or truncate 
    keys if not exactly 16 or 32 bytes as per...
        key:  0-15 bytes        pad to 16 bytes (128 bits)
        key: 17-31 bytes        pad to 32 bytes (256 bits)
        key:  > 32 bytes        truncate to 32 bytes
    Pads are nulls, 0x00.
    
    Because the key is typed, some will consider the key weak.  
    Stronger keys are possible if you modify this program.
    
    Nor does this program do a good job of managing memory and 
    ensuring keys written to memory get overwritten and erased.  
    Simply stated, it doesn't.  This program should NOT be used 
    on high value assets.
    
    This is EXPERIMENTAL software and intended for educational 
    purposes only.  To make experimentation less cumbersome, 
    hotSalsa is free for any use.      
    
    THIS PROGRAM IS PROVIDED WITHOUT WARRANTY OR GUARANTEE OF
    ANY KIND.  USE AT YOUR OWN RISK.  
    
    Larry Bugbee
    March 2007
    
"""

import pySalsa20, sys, time, os, getpass

#---------------------------------------------------------------
# constants

suffix = '.s20'

#---------------------------------------------------------------
# utilities

def overwritecheck(filename):
    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'
            sys.exit()
    
def savetofile(filename, content):
    overwritecheck(filename)
    f = open(filename, 'wb')
    f.write(content)
    f.close()
    
def loadfmfile(filename):
    f = open(filename, 'rb')
    content = f.read()
    f.close()
    return content

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():
    # sufficient delay on most processors to ensure a unique nonce?
    time.sleep(0.1)
    nonce = ('%f' % (time.time()))[-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
    

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

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

def main():
    # analyze file name suffix
    if len(sys.argv) > 1:
        filename = sys.argv[1]
        encrypt = not filename.endswith(suffix)
        if encrypt:                 # determined by file suffix    
            newfilename = filename+suffix
        else:
            newfilename = filename[:-len(suffix)]
    else:
        usage()
    
    # 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'
                sys.exit()
    key = key1632(pswd)         # pad or truncate as appropriate
    
    overwritecheck(newfilename)
    fi = open(filename, 'rb')
    fo = open(newfilename, 'wb')
    if encrypt:                 # determined by file suffix
        IV = iv8(makeNounce())
        fo.write(IV)            # front of file to have 8-byte IV
        ctx = pySalsa20.Salsa20(key, IV)
        t0 = time.time()
        while 1:
            chunk = fi.read(64000000)
            if not chunk: break
            ciphertext = ctx.encryptBytes(chunk)
            fo.write(ciphertext)
        print '%.3f seconds' % (time.time() - t0)
    else:
        IV = fi.read(8)         # get IV from front of file
        ctx = pySalsa20.Salsa20(key, IV)
        t0 = time.time()
        while 1:
            chunk = fi.read(64000000)
            if not chunk: break
            plaintext = ctx.decryptBytes(chunk)
            fo.write(plaintext)
        print '%.3f seconds' % (time.time() - t0)
    fi.close()
    fo.close()

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

if __name__ == '__main__':
    main()
    
#---------------------------------------------------------------
#---------------------------------------------------------------
#---------------------------------------------------------------