clear;
// random matrix M in Epm

random_M := function(m,p)
	Zp := RingOfIntegers(p);
	Zpm := RingOfIntegers(p^m);
	M := RandomMatrix(Zp,1,m);
	M := ChangeRing(M,Zpm);
	for i in [2..m] do
		c := RandomMatrix(RingOfIntegers(p^i),1,m);
		for j in [1..(i-1)] do
			c[1,j] := c[1,j]*p^(i-j);
		end for;
		M := VerticalJoin(M,ChangeRing(c,Zpm));
	end for;
	return M;
end function;


// Bringing matrix down to Epm

down_to_Epm := function(m,p,M)
	Zpm := RingOfIntegers(p^m);
	MM := ChangeRing(ChangeRing(M[1],RingOfIntegers(p)),Zpm);
	for i in [2..m] do
		MM := VerticalJoin(MM,ChangeRing(ChangeRing(M[i],RingOfIntegers(p^i)),Zpm));
	end for;
	return(MM);
end function;

// Bringing matrix up to FEpm

up_to_FEpm := function(m,p,M)
	Zpm := RingOfIntegers(p^m);
	for i in [1..m] do
		M[i] := p^(m-i)*M[i];
	end for;
	return(M);
end function;


// random matrix in H(M)

random_A := function(m,p,M)
	Zpm := RingOfIntegers(p^m);
	A := random_Z(m,p)*IdentityMatrix(Zpm,m);
	for i in [2..m] do
		A := A + random_Z(m,p)*M^i;
	end for;
	return down_to_Epm(m,p,A);
end function;


		
//The Attack (Algorithm 1)

attack := function(m,p,M,X,G1,G2)
	Zpm := RingOfIntegers(p^m);
	Zp := RingOfIntegers(p);
	//R<[x]> := PolynomialRing(Zpm,m^2);
	T := ZeroMatrix(Zpm,m^2,1);
	for i in [1..m] do
		for j in [1..m] do
			T := HorizontalJoin(T,Matrix(Zpm,m^2,1,ElementToSequence(up_to_FEpm(m,p,down_to_Epm(m,p,M^(i-1)*X*M^(j-1))))));
		end for;
	end for;	
	RemoveColumn(~T,1);
	g := Vector(Zpm,ElementToSequence(up_to_FEpm(m,p,G1)));
	//print T,g;
	lambda := Solution(Transpose(T),g);
	H := ZeroMatrix(Zpm,m,m);
	for i in [1..m] do
		for j in [1..m] do
			H := H + lambda[(i-1)*m+j]*M^(i-1)*G2*M^(j-1);
		end for;
	end for;
	return down_to_Epm(m,p,H);
end function;


//Setting m and p
m := 3;
p := 2;
Zpm := RingOfIntegers(p^m);
Zp := RingOfIntegers(p);
//Choosing random M and X
M := random_M(m,p);
X := random_M(m,p);


//Choosing A1,A2,B1,B2 in H(M) 
A1 := random_A(m,p,M);
A2 := random_A(m,p,M);
B1 := random_A(m,p,M);
B2 := random_A(m,p,M);

// Public info
G1 := down_to_Epm(m,p,A1*X*A2);
G2 := down_to_Epm(m,p,B1*X*B2);


//Applying the attack
time H := attack(m,p,M,X,G1,G2);

//Verifying the output with shared secret A1*G2*A2
if down_to_Epm(m,p,A1*G2*A2) eq H then
	print "success";
else
	print "failure";
end if;