from sage.all import *

class Encrypter(object):

    def __init__(self, matrixLength):
        super(Encrypter, self).__init__()
        self.numberSpace = ZZ
        self.matrixLength = matrixLength
        self.matrixSpace = MatrixSpace(self.numberSpace, matrixLength, matrixLength)
        self.generateSecretKey()

    def generateSecretKey(self):
        self.skVector = random_vector(self.numberSpace, self.matrixLength)
        print(self.skVector)

    def encrypt(self, opentext):
        skIndex = randrange(self.matrixLength)
        helpMatrix = self.matrixSpace()
        M = self.matrixSpace()
        DMatrix = self.matrixSpace()
        helpMatrix[skIndex] = self.skVector
        for i in range(self.matrixLength):
            if (i != skIndex):
                randVector = random_vector(self.numberSpace, self.matrixLength)
                helpMatrix[i] = randVector
        #if helpMatrix cannot be .transpose() cause determinant = 0
        while (helpMatrix.determinant() == 0):
            i = randrange(self.matrixLength)
            if (i != skIndex):
                helpMatrix[i] = random_vector(self.numberSpace, self.matrixLength)
        helpMatrix = helpMatrix.transpose()
        DMatrix[skIndex,skIndex] = opentext
        M = helpMatrix * DMatrix * helpMatrix.inverse()
        return M

    def decrypt(self, ciphertext):
        result = ciphertext * self.skVector
        for i in range(self.matrixLength):
            if (self.skVector[i] != 0):
                return result[i] / self.skVector[i]
        print("Error in decrypter the skVector have all zeros")
        return 0


matrixLength = 3
a = Encrypter(matrixLength)
c1 = a.encrypt(100)
print(c1)
c2 = a.encrypt(224)
print(c2)
cResult = c1 * c2
print(cResult)
a.decrypt(cResult)