#! /usr/bin/env python

# obelisk.py -- simple pseudo-one-time-pad en/decryption
# Copyright (C) 2008 Marshall T. Vandegrift <llasram@gmail.com>

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from optparse import OptionParser
from itertools import izip
import sys
import struct
import sha

class Obelisk(object):
    def __init__(self, salt, keyfile, digest=sha):
        self.counter = 0
        self.salt = salt
        self.keyfile = keyfile
        self.digest = digest
        self.keytext = ""

    def keystream(self, size):
        while len(self.keytext) < size:
            data = ""
            while len(data) < self.digest.digest_size:
                datum = self.keyfile.read(self.digest.digest_size - len(data))
                if datum == "":
                    self.keyfile.seek(0)
                    self.counter += 1
                data += datum
            salt = struct.pack(">I", self.counter) + self.salt
            self.keytext += self.digest.new(salt + data).digest()
        key, self.keytext = self.keytext[:size], self.keytext[size:]
        return key
        
    def decrypt(self, data):
        key = [ord(x) for x in self.keystream(len(data))]
        data = [ord(x) for x in data]
        result = [x ^ y for x, y in izip(key, data)]
        return ''.join(chr(x) for x in result)

def obelisk(salt, keyfile, infile, outfile):
    obelisk = Obelisk(salt, keyfile)
    data = infile.read(1024)
    while data:
        outfile.write(obelisk.decrypt(data))
        data = infile.read(1024)
    return True

if __name__ == '__main__':
    parser = OptionParser("usage: %prog SALT KEYFILE INFILE OUTFILE")
    (options, args) = parser.parse_args()
    if len(args) != 4:
        parser.error("incorrect number of arguments")
    (salt, keyfname, infname, outfname) = args
    try:
        keyfile = open(keyfname, 'rb')
    except Exception:
        parse.error("could not open key file '%s'" % (keyfname))
    try:
        infile = open(infname, 'rb')
    except Exception:
        parse.error("could not open input file '%s'" % (keyfname))
    try:
        outfile = open(outfname, 'wb')
    except Exception:
        parse.error("could not open output file '%s'" % (keyfname))

    if not obelisk(salt, keyfile, infile, outfile):
        sys.exit(1)

