1use base64::prelude::*;
22use once_cell::sync::Lazy;
23use rand::prelude::*;
24use rug::integer::Order;
25use rug::ops::Pow;
26use rug::Integer;
27use std::str;
28
29const VERSION: &str = "s";
30struct KctfParams {
31 pub modulus: Integer,
33 pub exponent: Integer,
35}
36
37impl KctfParams {
38 fn new() -> Self {
39 let big_num = Integer::from(2).pow(1279);
40 KctfParams {
41 modulus: big_num.clone() - 1,
42 exponent: big_num / 4,
43 }
44 }
45}
46
47static CUSTOM_BASE64: base64::engine::GeneralPurpose = base64::engine::GeneralPurpose::new(
48 &base64::alphabet::STANDARD,
49 base64::engine::GeneralPurposeConfig::new()
50 .with_decode_allow_trailing_bits(true)
51 .with_encode_padding(true)
52 .with_decode_padding_mode(base64::engine::DecodePaddingMode::Indifferent),
53);
54
55static KCTF_PARAMS: Lazy<KctfParams> = Lazy::new(|| KctfParams::new());
56
57#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
58pub enum KctfErrors {
59 UnknownVersion,
63 FormatError,
68 DecodeError,
74 LargeDifficulty,
76}
77
78#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
79pub struct KctfPow {
80 pub difficulty: u32,
82 pub value: Integer,
84}
85
86impl KctfPow {
87 pub fn from_difficulty_and_value(difficulty: u32, value: Integer) -> Self {
92 KctfPow { difficulty, value }
93 }
94
95 pub fn from_challenge(challenge: &str) -> Result<Self, KctfErrors> {
98 let mut challenge_parts = challenge.split('.');
99
100 if challenge_parts.next() != Some(VERSION) {
102 return Err(KctfErrors::UnknownVersion);
103 }
104
105 let challenge_parts: Vec<&str> = challenge_parts.collect();
106
107 if challenge_parts.len() != 2 {
109 return Err(KctfErrors::FormatError);
110 }
111
112 let decoded_parts: Vec<Vec<u8>> = challenge_parts
114 .into_iter()
115 .map(|x| CUSTOM_BASE64.decode(x).map_err(|_| KctfErrors::DecodeError))
116 .collect::<Result<_, KctfErrors>>()?;
117
118 let decoded_difficulty = &decoded_parts[0];
119 let decoded_value = &decoded_parts[1];
120
121 let difficulty: u32 = if decoded_difficulty.len() > 4 {
122 if (&decoded_difficulty[..decoded_difficulty.len() - 4])
123 .iter()
124 .any(|&x| x != 0)
125 {
126 return Err(KctfErrors::LargeDifficulty);
127 }
128 u32::from_be_bytes((&decoded_difficulty[..4]).try_into().unwrap())
129 } else {
130 let mut difficulty_array = [0; 4];
131 difficulty_array[4 - decoded_difficulty.len()..].copy_from_slice(&decoded_difficulty);
132 u32::from_be_bytes(difficulty_array)
133 };
134
135 Ok(KctfPow {
136 difficulty,
137 value: Integer::from_digits(decoded_value, Order::Msf),
138 })
139 }
140
141 pub fn solve(mut self) -> String {
143 for _ in 0..self.difficulty {
144 let _ = self
145 .value
146 .pow_mod_mut(&KCTF_PARAMS.exponent, &KCTF_PARAMS.modulus);
147 self.value ^= 1;
148 }
149
150 format!(
151 "{}.{}",
152 VERSION,
153 CUSTOM_BASE64.encode(self.value.to_digits(Order::Msf))
154 )
155 }
156
157 pub fn gen_challenge(difficulty: u32) -> Self {
159 let mut bytes: [u8; 16] = [0; 16];
160 thread_rng().fill(&mut bytes[..]);
161 KctfPow {
162 difficulty: difficulty,
163 value: Integer::from_digits(&bytes, Order::Msf),
164 }
165 }
166
167 pub fn serialize_challenge(&self) -> String {
169 format!(
170 "{}.{}.{}",
171 VERSION,
172 CUSTOM_BASE64.encode(self.difficulty.to_be_bytes()),
173 CUSTOM_BASE64.encode(self.value.to_digits(Order::Msf))
174 )
175 }
176
177 pub fn verify(&self, solution: &str) -> Result<bool, KctfErrors> {
179 let mut decoded_solution = decode_solution(solution)?;
180
181 for _ in 0..self.difficulty {
182 decoded_solution ^= 1;
183 let _ = decoded_solution.pow_mod_mut(&2.into(), &KCTF_PARAMS.modulus);
184 }
185
186 Ok(self.value == decoded_solution
187 || Integer::from(&KCTF_PARAMS.modulus - &self.value) == decoded_solution)
188 }
189 pub async fn async_solve(mut self) -> String {
191 for _ in 0..self.difficulty {
192 async {
193 let _ = self
194 .value
195 .pow_mod_mut(&KCTF_PARAMS.exponent, &KCTF_PARAMS.modulus);
196 self.value ^= 1;
197 }
198 .await
199 }
200
201 format!(
202 "{}.{}",
203 VERSION,
204 CUSTOM_BASE64.encode(self.value.to_digits(Order::Msf))
205 )
206 }
207
208 pub async fn async_verify(&self, solution: &str) -> Result<bool, KctfErrors> {
210 let mut decoded_solution = decode_solution(solution)?;
211
212 for _ in 0..self.difficulty {
213 async {
214 decoded_solution ^= 1;
215 let _ = decoded_solution.pow_mod_mut(&2.into(), &KCTF_PARAMS.modulus);
216 }
217 .await
218 }
219 Ok(self.value == decoded_solution
220 || Integer::from(&KCTF_PARAMS.modulus - &self.value) == decoded_solution)
221 }
222}
223
224pub fn decode_solution(solution: &str) -> Result<Integer, KctfErrors> {
228 let mut solution_parts = solution.split('.');
229 if solution_parts.next() != Some(VERSION) {
230 return Err(KctfErrors::UnknownVersion);
231 }
232
233 let decoded_solution = match solution_parts.next() {
234 Some(v) => Integer::from_digits(
235 &CUSTOM_BASE64
236 .decode(v)
237 .map_err(|_| KctfErrors::DecodeError)?,
238 Order::Msf,
239 ),
240
241 None => {
242 return Err(KctfErrors::FormatError);
243 }
244 };
245
246 if solution_parts.next().is_some() {
247 return Err(KctfErrors::FormatError);
248 }
249
250 Ok(decoded_solution)
251}