1
#!/usr/bin/env python3
2
# Solution for quantum-web-token CTF challenge
3

4
from pwn import * # Using pwntools for connection convenience
5
import base64
6
import json
7
import time
8

9
# JWT decoding function
10
def decode_jwt_payload_segment(payload_b64_str: str) -> dict:
11
    """
12
    Decodes the payload segment of a JWT token.
13
    The input is the base64url encoded payload string (middle part of JWT).
14
    """
15
    # Add padding if necessary for base64url decoding
16
    padded_payload_b64_str = payload_b64_str + '=' * (-len(payload_b64_str) % 4)
17
    payload_json_bytes = base64.urlsafe_b64decode(padded_payload_b64_str)
18
    return json.loads(payload_json_bytes.decode('utf-8'))
19

20
def solve(max_retries=5, initial_timeout=15, per_line_timeout=10, recv_timeout=10): # Increased timeouts
21
    for attempt in range(max_retries):
22
        log.info(f"Starting attempt {attempt + 1}/{max_retries}")
23
        conn = None  # Initialize conn to None
24
        try:
25
            # Set up pwntools context for quieter output or debugging
26
            # context.log_level = 'debug' # Uncomment for verbose pwntools output
27
            context.log_level = 'info'
28

29
            # conn = remote("localhost", 5001) # For local testing
30
            conn = remote("116.203.176.73", 5001, timeout=initial_timeout)
31
            log.info("Connected to server.")
32

33
            # Receive welcome message
34
            welcome_msg = conn.recvuntil(b"Sending qubits through quantum channel:\n", timeout=recv_timeout)
35
            log.info(f"Welcome message received: {welcome_msg.decode().strip()}")
36

37
            time.sleep(0.5) # Add a small delay after welcome, before server sends JWTs
38

39
            n_qubits = 32
40
            s_bits_master = [None] * n_qubits # Stores the known s_bits from Alice
41
            num_s_bits_known = 0
42
            current_round_xors = [None] * n_qubits # Stores XOR values for the current round (needs to be accessible after loop)
43

44

45
            round_num = 0
46
            # Loop until all of Alice's s_bits are known
47
            while num_s_bits_known < n_qubits:
48
                round_num += 1
49
                log.info(f"Starting Round {round_num}. Known s_bits: {num_s_bits_known}/{n_qubits}")
50

51
                # 1. Receive n_qubits qwt JWTs (individual qubit info)
52
                for i in range(n_qubits):
53
                    try:
54
                        jwt_line_bytes = conn.recvline(timeout=per_line_timeout)
55
                        if not jwt_line_bytes:
56
                            log.error(f"Connection closed prematurely while reading qwt JWT {i} in round {round_num}.")
57
                            raise EOFError("Connection closed by server")
58
                        jwt_line = jwt_line_bytes.strip().decode()
59

60
                        if jwt_line.count('.') != 2:
61
                            log.error(f"Received malformed JWT in round {round_num}, JWT {i}: {jwt_line}")
62
                            raise ValueError("Malformed JWT received")
63

64
                        payload_b64_str = jwt_line.split('.')[1]
65
                        payload = decode_jwt_payload_segment(payload_b64_str)
66
                        current_round_xors[i] = payload['xor']
67
                    except Exception as e:
68
                        log.error(f"Error processing qwt JWT {i} in round {round_num}: {e}")
69
                        problematic_line = jwt_line_bytes if 'jwt_line_bytes' in locals() and jwt_line_bytes is not None else b'N/A'
70
                        log.error(f"Problematic JWT line (if any): {problematic_line.decode(errors='ignore')}")
71
                        raise
72
                log.info(f"Round {round_num}: Received all {n_qubits} qwt JWTs.")
73

74
                # 2. Receive p_key JWT (sifted key info)
75
                try:
76
                    pkey_jwt_line_bytes = conn.recvline(timeout=per_line_timeout)
77
                    if not pkey_jwt_line_bytes:
78
                        log.error(f"Connection closed prematurely while reading pkey JWT in round {round_num}.")
79
                        raise EOFError("Connection closed by server")
80
                    pkey_jwt_line = pkey_jwt_line_bytes.strip().decode()
81

82
                    if pkey_jwt_line.count('.') != 2:
83
                        log.error(f"Received malformed pkey JWT in round {round_num}: {pkey_jwt_line}")
84
                        raise ValueError("Malformed pkey JWT received")
85

86
                    pkey_payload_b64_str = pkey_jwt_line.split('.')[1]
87
                    pkey_payload = decode_jwt_payload_segment(pkey_payload_b64_str)
88
                except Exception as e:
89
                    log.error(f"Error processing pkey JWT in round {round_num}: {e}")
90
                    problematic_pkey_line = pkey_jwt_line_bytes if 'pkey_jwt_line_bytes' in locals() and pkey_jwt_line_bytes is not None else b'N/A'
91
                    log.error(f"Problematic pkey JWT line (if any): {problematic_pkey_line.decode(errors='ignore')}")
92
                    raise
93

94
                new_bits_learned_this_round = 0
95
                for idx_str, bit_val in pkey_payload.items():
96
                    idx = int(idx_str)
97
                    if s_bits_master[idx] is None:
98
                        s_bits_master[idx] = bit_val
99
                        num_s_bits_known += 1
100
                        new_bits_learned_this_round +=1
101

102
                log.info(f"Round {round_num}: Learned {new_bits_learned_this_round} new s_bits. Total known: {num_s_bits_known}/{n_qubits}")
103

104
                # 3. Decide whether to continue to another round or stop
105
                try:
106
                    prompt = conn.recvuntil(b"Do you want somemore tokens? (y/n): ", timeout=recv_timeout).decode()
107
                except Exception as e:
108
                    log.error(f"Error receiving 'more tokens' prompt in round {round_num}: {e}")
109
                    raise
110

111
                if num_s_bits_known == n_qubits:
112
                    log.info("All s_bits are known. Sending 'n' to stop receiving tokens.")
113
                    conn.sendline(b"n")
114
                    break
115
                else:
116
                    conn.sendline(b"y")
117

118
            log.info("All s_bits collected. Calculating the final r_bits (Bob's measured bits).")
119

120
            r_bits_final = [None] * n_qubits
121
            for i in range(n_qubits):
122
                if s_bits_master[i] is None:
123
                    log.error(f"Critical error: s_bits_master[{i}] is None even after loop completion.")
124
                    raise ValueError(f"s_bits_master[{i}] is None")
125
                if current_round_xors[i] is None: # Should not happen if logic is correct
126
                    log.error(f"Critical error: current_round_xors[{i}] is None before final calculation.")
127
                    raise ValueError(f"current_round_xors[{i}] is None")
128
                r_bits_final[i] = s_bits_master[i] ^ current_round_xors[i]
129

130
            try:
131
                guess_prompt = conn.recvuntil(b"Guess Bob's qubits (comma as delimeter): ", timeout=recv_timeout).decode()
132
            except Exception as e:
133
                log.error(f"Error receiving 'guess' prompt: {e}")
134
                raise
135

136
            guess_str = ",".join(map(str, r_bits_final))
137
            log.info(f"Sending guess: {guess_str}")
138
            conn.sendline(guess_str.encode())
139

140
            try:
141
                result_bytes = conn.recvall(timeout=recv_timeout + 10)
142
                result = result_bytes.decode(errors='ignore')
143
                log.info(f"Server response:\n{result}")
144
                if "UMCS{" in result:
145
                    log.success("Flag captured!")
146
                    print(f"FLAG: {result.strip()}")
147
                elif "worst guess ever" in result:
148
                    log.error("Guess was incorrect.")
149
                else:
150
                    log.warning("Flag not found in the response, and no explicit failure message.")
151
                return
152
            except Exception as e:
153
                log.error(f"Error receiving final result: {e}")
154
                raise
155

156
        except EOFError as e:
157
            log.error(f"EOFError during attempt {attempt + 1}: {e}. Server closed connection.")
158
        except PwnlibException as e:
159
            log.error(f"PwnlibException during attempt {attempt + 1}: {e}")
160
        except ValueError as e:
161
            log.error(f"ValueError during attempt {attempt + 1}: {e}")
162
        except Exception as e:
163
            log.error(f"An unexpected error occurred during attempt {attempt + 1}: {type(e).__name__} - {e}")
164

165
        if conn:
166
            conn.close()
167

168
        if attempt < max_retries - 1:
169
            log.info(f"Retrying in 5 seconds...")
170
            time.sleep(5)
171
        else:
172
            log.error("All retries failed.")
173
            break
174

175
    log.info("Solver finished.")
176

177
if __name__ == '__main__':
178
    solve()