Dragon CTF 2020 | Bit Flip 1 Solution

drgns2020_bitflip1.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from pwn import *
import subprocess
from Crypto.Util.number import bytes_to_long, long_to_bytes
from Crypto.Cipher import AES
import hashlib
import os
import base64
from gmpy2 import is_prime

context.update(arch='i386')
exe = './task.py'

host = args.HOST or 'bitflip1.hackable.software'
port = int(args.PORT or 1337)

### COPIED
class Rng:
  def __init__(self, seed):
    self.seed = seed
    self.generated = b""
    self.num = 0

  def more_bytes(self):
    self.generated += hashlib.sha256(self.seed).digest()
    self.seed = long_to_bytes(bytes_to_long(self.seed) + 1, 32)
    self.num += 256


  def getbits(self, num=64):
    while (self.num < num):
      self.more_bytes()
    x = bytes_to_long(self.generated)
    self.num -= num
    self.generated = b""
    if self.num > 0:
      self.generated = long_to_bytes(x >> num, self.num // 8)
    return x & ((1 << num) - 1)


class DiffieHellman:
  def gen_prime(self):
    prime = self.rng.getbits(512)
    iter = 0
    while not is_prime(prime):
      iter += 1
      prime = self.rng.getbits(512)
    return prime

  def __init__(self, seed, prime=None):
    self.rng = Rng(seed)

    if prime is None:
      prime = self.gen_prime()

    # p = self.rng.getbits(512) until prime
    self.prime = prime

    # a = self.rng.getbits(64)
    self.my_secret = self.rng.getbits()

    # A = g^a mod p, g=5
    self.my_number = pow(5, self.my_secret, prime)
    self.shared = 1336

  def set_other(self, x):
    self.shared ^= pow(x, self.my_secret, self.prime)
######################


def local(argv=[], *a, **kw):
    '''Execute the target binary locally'''
    if args.GDB:
        return gdb.debug([exe] + argv, gdbscript=gdbscript, *a, **kw)
    else:
        return process([exe] + argv, *a, **kw)

def remote(argv=[], *a, **kw):
    '''Connect to the process on the remote host'''
    io = connect(host, port)
    if args.GDB:
        gdb.attach(io, gdbscript=gdbscript)
    return io

def start(argv=[], *a, **kw):
    '''Start the exploit against the target.'''
    if args.LOCAL:
        return local(argv, *a, **kw)
    else:
        return remote(argv, *a, **kw)

# ./exploit.py GDB
gdbscript = '''
continue
'''.format(**locals())

# send xor string
def send_info(flip):
    io.recvuntil('str:')
    io.recvline()
    io.sendline(base64.b64encode(flip))
    x = io.recvline().strip().decode()
    #print(x)
    i = int(x.split(" ")[2])
    #print("FLIP:", flip)
    #print("ITERATIONS:", i)
    b = io.recvline().strip().decode()
    #print("BOB:", b)
    iv = io.recvline().strip().decode()
    #print("IV:", iv)
    flag = io.recvline().strip().decode()
    #print("FLAG:", flag)
    #print("SEED:", io.recvline().strip().decode())
    return i, b.split(" ")[2], iv, flag

# sets all known bits 1, or sets all known bits to 0 and flips the next unknown bit to add 2
def get_flip(k, r, unknown):
    mask = 0
    for i in range(1, k+1):
        mask |= 1 << i

    if unknown:
        flip = mask & r
        flip |= (1 << k+1)
    else:
        flip = mask & (~r)

    flip = int.to_bytes(flip, 32, 'big')
    return flip

io = start()

# pow
if not args.LOCAL:
    io.recvuntil("Work: ")
    proof = io.recvline()

    print(proof.decode('utf-8'))
    token = subprocess.check_output(proof.decode('utf-8'), shell=True)
    print(token)
    io.send(token)

j = 1
knowledge = 0
result = 0
while True:
    # get first mask that sets known bits to 1
    flip = get_flip(knowledge, result, unknown=False)
    # first time through to get i
    iterations = send_info(flip)[0]
    # get second mask that sets known bits to 0 and unknown bit to 1
    flip = get_flip(knowledge, result, unknown=True)
    # second time through, comparing i
    iterations2 = send_info(flip)[0]

    # THIS BIT POSITION IS 0,
    if (iterations - 1 == iterations2):
        knowledge += 1
        #print(knowledge, "bit is 0!")
    # THIS BIT POSITION IS 1
    else:
        knowledge += 1
        #print(knowledge, "bit is 1!")
        result |= (1 << j)
    j += 1
    if j > 128:
        break


# test out our derived seed, not sure if last bit is 1 or 0 though
def test(flip):
    iterations, bob, iv, flag = send_info(flip)
    alice_seed = long_to_bytes(result, 32)
    #print("SEEEEEEED:", alice_seed)
    alice = DiffieHellman(alice_seed)
    alice.set_other(int(bob))
    iv = base64.b64decode(iv)

    #print("iv:", iv)

    cipher = AES.new(long_to_bytes(alice.shared, 16)[:16], AES.MODE_CBC, IV=iv)
    f = cipher.decrypt(base64.b64decode(flag))
    print(f)

flip = b'\x00' * 32
test(flip)
flip = b'\x00' * 31 + b'\x01'
test(flip)