# -*- coding: windows-1251 -*- 
"""Encryptor supports encryption/decryption with
different algorithms. Type parameter --help to see more info"""

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

import sys
import getopt
from copy import copy

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

def print_help():
    """This function prints help for encryptor. To call it
        you cat type parameter -h or --help"""
    print   "\n--help [-h]\t\t\t\t- view this help \
             \n--encrypt [-e] <algorythm>\t\t- file encryption \
             \n--decrypt [-d] <algorythm>\t\t- file decryption \
             \n--source [-s] <filename>\t\t- source file \
             \n\nAvailable algorythms: \
             \nsm\t - simple matrix \
             \nkm\t - simple matrix with key"

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

def simple_matrix_cryptor(source, crypt_mode):
    """This function encrypts contents of source file with
        simple matrix method"""
    try:
        filename = source.split('.')[0].split('\\')[1]
        dir = source.split('.')[0].split('\\')[0]
    except Exception:
        try: 
            filename = source.split('.')[0]
            dir = ""
        except Exception:
            filename = source
    if dir != "":
        dir += "\\"
    
    print "Simple matrix encryption [", source, "]"
    print "Reading source . . ."
    
    source_data = ""
    try:
        fsource = open(source, "rb")
        source_data = fsource.read()
        fsource.close()
    except IOError:
        print "File ", source, " not found"
        sys.exit()
    
    dimention = 0
    if source_data != "":
        #    what dimention
        while dimention**2 < source_data.__len__():
            dimention += 1

        #    making 2D array
        matrix = [copy('\x00') for i in range(dimention)]
        matrix = [copy(matrix) for i in range(dimention)]
        
        print "Making 2D array . . ."
        #    filling array with given data
        for i in range(dimention):
            for j in range(dimention):
                try:
                    if crypt_mode == "encrypt":
                        matrix[i][j] = source_data[j + dimention * i]
                    elif crypt_mode == "decrypt":
                        matrix[i][j] = source_data[i + dimention * j]
                except Exception:
                    pass

        print "Encrypting data . . ."
        #    making encrypted data
        result = ''
        for i in range(dimention):
            for j in range(dimention):
                if crypt_mode == "encrypt":
                        result += matrix[j][i].__str__()
                elif crypt_mode == "decrypt":
                    if matrix[i][j] != '\x00':
                        result += matrix[i][j].__str__()

        #    writing result to file
        try:
            if crypt_mode == "encrypt":
                fresult = open(dir + filename + '.smc', 'wb')
            elif crypt_mode == "decrypt":
                fresult = open(dir + filename + '.smd', 'wb')
            fresult.write(result)
            fresult.close()
        except Exception:
            print "Can't open file ", dir, filename, '.smc for writing'
            sys.exit()
        
        print "Done." 

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

def key_matrix_cryptor(source):
    print "simple matrix encryption with key"
    
#------------------------------------------------------------------------------

if __name__ == "__main__":
    try:
        del(sys.argv[0:1])
        opts, args = getopt.getopt(sys.argv, 'he:d:s:', ['help', 'encryption=', 'decryption=', 'source='])
        #    parsing arguments
        for opt, arg in opts:
            
            #    executing help
            if opt in ('-h', '--help'):
                print_help()
                
            #    encrypting
            elif opt in('-e', '--encrypt'):
                try:
                    source = opts[1][1]
                except Exception:
                    print "No parameter with source file"
                    sys.exit()
                if arg == 'sm':
                    simple_matrix_cryptor(source, "encrypt")
                elif arg == 'km':
                    key_matrix_cryptor(source)
            
            #    decrypting
            elif opt in('-d', '--decrypt'):
                try:
                    source = opts[1][1]
                except Exception:
                    print "No parameter with source file"
                    sys.exit()
                if arg == 'sm':
                    simple_matrix_cryptor(source, "decrypt")
                    
    #    wrong arguments
    except getopt.GetoptError:
        sys.exit()