SpamAndHex CTF 2020 - OTS task

solve.rs

use std::net::TcpStream;
use std::io::{BufRead, BufReader, Write, BufWriter};
use regex::Regex;

fn wrap(data: &[u8]) -> Vec<u8> {
    let mut vec = data.to_vec();
    assert!(data.len() <= 128 - 16);
    vec.resize(128 - 16, 0u8);
    let dig = md5::compute(&vec);
    vec.extend(dig.iter());

    vec
}

fn hash_iter(data: &[u8], iterations: usize) -> Vec<u8> {
    let mut vec = data.to_vec();
    for _ in 0..iterations {
        vec = md5::compute(&vec).to_vec();
    }
    vec
}

fn verify(pub_key: &[u8], msg: &[u8], signature: &[u8]) -> bool {
    let wrapped = wrap(msg);
    assert!(signature.len() == 16 * 128 && signature.len() == pub_key.len());
    let decoded_pub_key = signature
        .chunks_exact(16)
        .zip(wrapped.iter())
        .map(|(chunk, iters)| hash_iter(chunk, *iters as usize))
        .flatten()
        .collect::<Vec<_>>();

    pub_key == decoded_pub_key.as_slice()
}

fn solve(pub_key: Vec<u8>, msg: String, signature: Vec<u8>) -> Result<(String, Vec<u8>), &'static str> {
    let flag = b"flag";
    let wrapped_msg = wrap(msg.as_bytes());

    let blocked_range = 5..9; // "My fa[vori]te num..."
    let mut flag_msg = msg.as_bytes().to_vec();
    flag_msg[blocked_range.clone()].copy_from_slice(flag);

    let mut comb = 0usize;
    loop {
        if comb % 10000 == 0 {
            println!("{}", comb);
        }
        let mut curcomb = comb;
        let mut curmsg = flag_msg.clone();
        for (i, x) in curmsg.iter_mut().enumerate() {
            if blocked_range.contains(&i) {
                continue;
            }
            if curcomb == 0 {
                break;
            }
            let m = *x - 20 + 1;
            let delta = curcomb % m as usize; 
            *x -= delta as u8;
            curcomb /= m as usize;
        }
        if curcomb > 0 {
            return Err("No solution");
        }
        let cur_wrapped = wrap(&curmsg);
        if cur_wrapped[(128 - 16)..].iter().zip(&wrapped_msg[(128 - 16)..]).all(|(a, b)| a <= b) {
            let new_msg = String::from_utf8_lossy(&cur_wrapped[..msg.len()]);
            let mut new_signature = signature.clone();
            for ((x, y), chunk) in cur_wrapped.iter().zip(&wrapped_msg).zip(new_signature.chunks_exact_mut(16)) {
                assert!(x <= y);
                let iters = y - x;
                chunk.copy_from_slice(&hash_iter(chunk, iters as usize));
            }
            assert!(verify(&pub_key, new_msg.as_bytes(), &new_signature));
            return Ok((new_msg.to_string(), new_signature))
        } else {
            comb += 1;
        }
    }
}

fn main() -> std::io::Result<()> {
    let stream = TcpStream::connect("34.89.64.81:1337")?;
    let mut reader = BufReader::new(stream.try_clone().expect("Stream clone failed"));
    let mut writer = BufWriter::new(stream);

    let mut buf = String::new();
    while !buf.ends_with("Enter the message:\n") {
        reader.read_line(&mut buf)?;
    }

    let re = Regex::new(r#"pub_key = ([^\s]+)\nsign\("([^"]+)"\) = ([^\s]+)"#).unwrap();
    let cap = re.captures(&buf).unwrap();
    let pub_key = hex::decode(cap.get(1).unwrap().as_str()).unwrap();
    let msg = cap.get(2).unwrap().as_str().to_string();
    let signature = hex::decode(cap.get(3).unwrap().as_str()).unwrap();

    let (new_msg, solution) = solve(pub_key, msg, signature).expect("Solution not found");
    println!("Prepared message: \"{}\"", new_msg);

    // Send the message
    writeln!(writer, "{}", new_msg)?;
    writer.flush()?;

    buf.clear();
    reader.read_line(&mut buf)?;
    println!("Received: [{}]", buf.trim()); // "Enter the signature:"

    // Send the signature
    writeln!(writer, "{}", hex::encode(&solution))?;
    writer.flush()?;

    buf.clear();
    reader.read_line(&mut buf)?;
    println!("Received: [{}]", buf.trim()); // The flag

    Ok(())
}