[PokaCTF] Igo-Genni

from sage.all import *  
import os, pickle, binascii
from Crypto.Hash import SHA256
 
## Class of descending l-isogeny chains
class Chain:
	def __init__(self,osidh,L_j=[]):
 
		self.osidh=osidh
		if len(L_j)>0:
			self.L_j=L_j
		else:
			self.L_j=[]
			if osidh.d_K==-3:
				self.L_j.append(osidh.F(0))
			else:
				self.L_j.append(osidh.F(1728))
 
			phi_l=osidh.L_phi[0]
			Fz=osidh.Fz
 
			if osidh.n>=1:
				# Choice of the second j-invariant
				L_eval=[]
				for g in phi_l:
					L_eval.append(g(self.L_j[0]))
				f=Fz(L_eval)
				L_roots=f.roots(multiplicities=False)
				m=len(L_roots)
				i=randint(0,m-1)
				while L_roots[i]==self.L_j[0]:
					i=randint(0,m-1)
				self.L_j.append(L_roots[i])
 
				# Choice of the other j-invariants
				k=2
				while k<=osidh.n:
					L_eval=[]
					for g in phi_l:
						L_eval.append(g(self.L_j[k-1]))
					f=Fz(L_eval)
					L_roots=f.roots(multiplicities=False)
					m=len(L_roots)
					i=randint(0,m-1)
					while L_roots[i]==self.L_j[k-2]:
						i=randint(0,m-1)
					self.L_j.append(L_roots[i])
					k+=1
 
## Class of horizontal isogeny chains
class Chain_hor:
	def __init__(self,osidh,ind_q,j_center,L_plus,L_minus):
 
		self.osidh=osidh
		self.ind_q=ind_q
		self.j_center=j_center
		self.L_plus=L_plus
		self.L_minus=L_minus
 
	def action_step(self,ind_q,j,e):
 
		phi_q1=self.osidh.L_phi[self.ind_q+1] 
		phi_q2=self.osidh.L_phi[ind_q+1] 
		Fz=self.osidh.Fz
 
		L_plus=[j]
		L_minus=[j]
 
		if e>=0:
			k=1
			while k<=e:
				L_eval=[]
				for g in phi_q1:
					L_eval.append(g(L_plus[-1]))
				f_q1=Fz(L_eval)
				L_eval=[]
				for g in phi_q2:
					L_eval.append(g(self.L_plus[k-1]))
				f_q2=Fz(L_eval)
				f=gcd(f_q1,f_q2)
				L_roots=f.roots(multiplicities=False)
				L_plus.append(L_roots[0])
				k+=1
		else:
			k=1
			while k<=-e:
				L_eval=[]
				for g in phi_q1:
					L_eval.append(g(L_minus[-1]))
				f_q1=Fz(L_eval)
				L_eval=[]
				for g in phi_q2:
					L_eval.append(g(self.L_minus[k-1]))
				f_q2=Fz(L_eval)
				f=gcd(f_q1,f_q2)
				L_roots=f.roots(multiplicities=False)
				L_minus.append(L_roots[0])
				k+=1
		return Chain_hor(self.osidh,self.ind_q,j,L_plus[1::],L_minus[1::])
 
N = 28      
T = 10      
L = 2       
R = 3       
D_K = -4    
 
def kdf_from_j(j):
    return SHA256.new(str(j).encode()).digest()
 
def xor_bytes(a: bytes, b: bytes) -> bytes:
    n = min(len(a), len(b))
    return bytes(x ^ y for x, y in zip(a[:n], b[:n]))
 
osidh = OSIDH(N, T, L, R, D_K)              
pub_chain = Chain(osidh)                      
 
L_exp_A = [randint(-osidh.r, osidh.r) for _ in range(osidh.t)]
L_exp_B = [randint(-osidh.r, osidh.r) for _ in range(osidh.t)]
 
chain_A = pub_chain.action(L_exp_A)           
chain_B = pub_chain.action(L_exp_B)
 
A_hor = []
B_hor = []
for j in range(osidh.t):
 
    # Alice’s horizontals
    ch = chain_A
    L_plus = []
    for _ in range(osidh.r):
        ch = ch.action_prime(osidh.L_mfq[j], j)
        L_plus.append(ch.L_j[-1])
    ch = chain_A
    L_minus = []
    for _ in range(osidh.r):
        ch = ch.action_prime(osidh.L_mfq_inv[j], j)
        L_minus.append(ch.L_j[-1])
    A_hor.append(Chain_hor(osidh, j, chain_A.L_j[-1], L_plus, L_minus))
 
    # Bob’s horizontals
    ch = chain_B
    L_plus = []
    for _ in range(osidh.r):
        ch = ch.action_prime(osidh.L_mfq[j], j)
        L_plus.append(ch.L_j[-1])
    ch = chain_B
    L_minus = []
    for _ in range(osidh.r):
        ch = ch.action_prime(osidh.L_mfq_inv[j], j)
        L_minus.append(ch.L_j[-1])
    B_hor.append(Chain_hor(osidh, j, chain_B.L_j[-1], L_plus, L_minus))
 
shared_j = chain_B.action(L_exp_A).L_j[-1]
 
FLAG_PATH = os.environ.get("FLAG_PATH", "flag.txt")
flag = open(FLAG_PATH, "rb").read().strip()
 
key = kdf_from_j(shared_j)
ct = xor_bytes(flag, key)
 
with open("osidh.pkl", "wb") as f:
    pickle.dump(osidh, f)
with open("pub_chain.pkl", "wb") as f:
    pickle.dump(pub_chain, f)
with open("A_hor.pkl", "wb") as f:
    pickle.dump(A_hor, f)
with open("B_hor.pkl", "wb") as f:
    pickle.dump(B_hor, f)
with open("flag.enc", "w") as f:
    f.write(binascii.hexlify(ct).decode())