#!/usr/bin/env python3
"""
XTS Forgery Attack on AES-XTS disk encryption service.

XTS mode: CT_i = E_K1(PT_i XOR T_i) XOR T_i
where T_i = E_K2(tweak_i), tweak_i = block_index as 16-byte little-endian.

The debug ECB oracle lets us compute:
  - E_K1(anything) using oracle slot 1
  - E_K2(anything) using oracle slot 2

So to forge block i with desired plaintext PT:
  1. T_i = E_K2(i.to_bytes(16, 'little'))        [oracle slot 2]
  2. inner = PT XOR T_i
  3. forged_CT = E_K1(inner) XOR T_i              [oracle slot 1]

Goal: make alice (block 71) have uid=0000 and gid=0000 so uid==0, gid==0.
Block 71 original: b'\nalice:x:1000:10'
Block 71 target:   b'\nalice:x:0000:00'
(block 72 stays b'00:Alice Example' -> gid becomes 0000 ✓)
"""

from pwn import remote

def xor_bytes(a, b):
    return bytes(x ^ y for x, y in zip(a, b))

TARGET_BLOCK_INDEX = 71
TARGET_PT = b'\nalice:x:0000:00'
assert len(TARGET_PT) == 16

tweak = TARGET_BLOCK_INDEX.to_bytes(16, 'little')

print("[*] Connecting...")
r = remote('diskenc-5df22768.camp09.c4mp.site', 1337, ssl=True)

# Read intro
intro = r.recvuntil(b"I can encrypt two blocks under ECB.\n")
print("[*] Server:", intro[-80:].decode(errors='replace'))

# We need:
#   oracle1_input = TARGET_PT XOR T_71  (but T_71 = E_K2(tweak), computed by oracle2)
# Problem: we need T_71 first to compute oracle1_input, but we only get one query.
# Solution: send tweak as oracle2 input first, get T_71, then... 
# BUT we have to send BOTH inputs in one line!
#
# Alternative approach: send a dummy for oracle1, get T_71 from oracle2,
# then compute what we need... but that's two rounds.
#
# Actually re-reading: we send one line with TWO hex blocks.
# oracle1 encrypts block1 with K1, oracle2 encrypts block2 with K2.
# We need T_71 = E_K2(tweak) to compute oracle1 input.
# 
# Two-step approach:
#   Round 1: send dummy|tweak -> get T_71 from oracle2 output
#   But we only get one oracle query!
#
# Wait - we can use a two-round trick:
#   Step 1: query oracle2 with tweak to get T_71
#           query oracle1 with zeros (dummy)
#   Step 2: compute inner = TARGET_PT XOR T_71
#           But we already used our oracle query...
#
# We need to do this in ONE shot. Let's think differently:
# Send: oracle1_input=tweak, oracle2_input=tweak
# Get: E_K1(tweak) and T_71=E_K2(tweak)
# Then: inner = TARGET_PT XOR T_71
# But now we need E_K1(inner), not E_K1(tweak)...
#
# SOLUTION: We can reconnect! Each connection is a fresh session.
# Connection 1: get T_71 = E_K2(tweak) by sending "00..00 <tweak_hex>"
# Connection 2: compute inner = TARGET_PT XOR T_71, send "<inner_hex> <tweak_hex>"
#               -> get E_K1(inner) and T_71 again
#               -> forged_CT = E_K1(inner) XOR T_71

r.close()

# --- Connection 1: get T_71 ---
print("[*] Connection 1: getting T_71 = E_K2(tweak_71)...")
r = remote('diskenc-5df22768.camp09.c4mp.site', 1337, ssl=True)
r.recvuntil(b"I can encrypt two blocks under ECB.\n")

dummy = b'\x00' * 16
query1 = dummy.hex() + " " + tweak.hex()
print(f"[*] Sending ECB query: {query1}")
r.sendline(query1.encode())

ecb_response = r.recvuntil(b"\n").strip()
print(f"[*] ECB response: {ecb_response}")
parts = ecb_response.split()
# parts[0] = E_K1(dummy), parts[1] = E_K2(tweak) = T_71
T_71 = bytes.fromhex(parts[1].decode())
print(f"[*] T_71 = {T_71.hex()}")
r.close()

# --- Connection 2: forge the block ---
print("\n[*] Connection 2: forging block 71...")
r = remote('diskenc-5df22768.camp09.c4mp.site', 1337, ssl=True)
r.recvuntil(b"I can encrypt two blocks under ECB.\n")

inner = xor_bytes(TARGET_PT, T_71)
query2 = inner.hex() + " " + tweak.hex()
print(f"[*] Sending ECB query: {query2}")
r.sendline(query2.encode())

ecb_response2 = r.recvuntil(b"\n").strip()
print(f"[*] ECB response: {ecb_response2}")
parts2 = ecb_response2.split()
E_K1_inner = bytes.fromhex(parts2[0].decode())
T_71_again = bytes.fromhex(parts2[1].decode())
assert T_71 == T_71_again, "T_71 mismatch!"

forged_CT = xor_bytes(E_K1_inner, T_71)
print(f"[*] Forged ciphertext for block {TARGET_BLOCK_INDEX}: {forged_CT.hex()}")

# --- Submit forged block ---
r.recvuntil(b"give me number and contents (but do not change usernames).\n")
block_update = f"{TARGET_BLOCK_INDEX} {forged_CT.hex()}"
print(f"[*] Submitting: {block_update}")
r.sendline(block_update.encode())

# Read response
response = r.recvall(timeout=5)
print("\n[*] Server response:")
print(response.decode(errors='replace'))
r.close()