exploit.py

from sm2 import P, A, B, GX, GY, N
from client_handshake import ClientHandshake

import random
import hashlib
import socket
import math
from typing import List, Tuple

from Crypto.Util.number import long_to_bytes, bytes_to_long
from Crypto.Cipher import AES

from sage.all import GF, EllipticCurve, CRT_list, prime_factors, is_prime
from sage.schemes.elliptic_curves.ell_point import EllipticCurvePoint_finite_field


SERVER_ADDRESS = "127.0.0.1"
SERVER_PORT = 8080

def is_key_valid(key: bytes, nonce: bytes, encrypted_signature: bytes) -> bool:
	"""
	Try to decrypt the encrypted signature with the given key and nonce.
	check if the key was correct
	"""
	cipher = AES.new(key, AES.MODE_CTR, nonce=b"", initial_value=nonce)
	decrypted_data = cipher.decrypt(encrypted_signature)
	return decrypted_data.startswith(b"LlsServerHello:")


def point_to_bytes(point: EllipticCurvePoint_finite_field) -> bytes:
	"""
	Convert a point to bytes: 32 bytes for x and 32 bytes for y.
	"""
	x_bytes = point.x().to_bytes(byteorder="big")
	y_bytes = point.y().to_bytes(byteorder="big")
	assert len(x_bytes) == 32
	assert len(y_bytes) == 32

	return x_bytes + y_bytes


def prime_factors_in_range(n: int, r: int) -> List[int]:
	"""
	Find all prime factors of n in the range [1, r].
	"""
	factors = []
	for i in range(2, r + 1):
		if n % i == 0 and is_prime(i):
			factors.append(i)
	
	return factors


def key_from_point(point: EllipticCurvePoint_finite_field, i: int) -> bytes:
	point_guess = point * i
	# This is like that because of the go implemntation of point to string
	key_string = f"X: {int(point_guess.x())}, Y: {int(point_guess.y())}".encode()
	return hashlib.sha256(key_string).digest()


def get_reminder(G_low_order: EllipticCurvePoint_finite_field, prime: int) -> Tuple[int, int]:
	# Client private key doesn't matter, we just need to get the encrypted_signature and nonce
	client_private_key = bytes_to_long(bytes([0xff] * 32))
	client_handshake = ClientHandshake(client_private_key, SERVER_ADDRESS, SERVER_PORT)
	encrypted_signature, nonce = client_handshake.get_encryption_given_G(point_to_bytes(G_low_order))

	for i in range(1, prime):
		k = key_from_point(G_low_order, i)
		if is_key_valid(k, nonce, encrypted_signature):
			print(f"Found key: {k.hex()}, {i} mod {prime}")
			return i, prime
	
	return 0, prime


def find_reminders_of_private_key(already_seen_primes: List[int]) -> List[Tuple[int, int]]:
	if not already_seen_primes:
		# A cool rand_b that has many small primes and easily factored - nice to start with
		rand_b = 47697949902441150779834732026565258987095441633064994778645432741922495526433
	else:
		rand_b = random.randint(1, P)
	
	E = EllipticCurve(GF(P), [A, rand_b])
	order = E.order()
	print(f"Order of curve: {order}")
	factors = prime_factors_in_range(order, 2000)
	factors = list(set(factors) - set(already_seen_primes))
	print(f"Small new factors of order of curve: {factors}")

	G = (GX, GY)
	reminders = []
	for factor in factors:
		G_low_order = E.gen(0) * int(order / factor)
		reminders.append(get_reminder(G_low_order, factor))
	
	return reminders


def main() -> None:
	# current_crt =  [(1, 17), (1, 2), (1, 3), (3, 5), (4, 7), (4, 13), (7, 29), (796, 1451), (331, 653), (955, 1223), (32, 41), (33, 257), (7, 499), (40, 61), (0, 11), (988, 1439), (25, 673), (24, 463), (78, 191), (17, 47), (10, 23), (378, 599), (1329, 1667), (35, 127), (27, 67), (315, 1109), (843, 941), (71, 113), (45, 71), (19, 73), (207, 307), (965, 1051), (1697, 1979), (3, 19), (531, 821), (31, 43), (513, 997), (6, 31), (222, 839), (134, 163)]	
	current_crt = []
	while len(current_crt) < 40:
		already_seen_primes = [_[1] for _ in current_crt]
		current_crt += find_reminders_of_private_key(already_seen_primes)
		print(f"Found {len(current_crt)} crt pairs:", current_crt)
	
	server_private_key = CRT_list([_[0] for _ in current_crt], [_[1] for _ in current_crt])
	server_private_key_bytes = long_to_bytes(server_private_key)
	print(f"Found private key: '{server_private_key_bytes.hex()}', " \
		  f"of length: {len(server_private_key_bytes)}, as int: {server_private_key}")


if __name__ == "__main__":
	main()