Joints 2019 Qual - Crypto

#!/usr/bin/env python

from Crypto.Cipher import AES
import os
import sys
import random
from base64 import *

class Unbuffered(object):
  def __init__(self, stream):
    self.stream = stream
  def write(self, data):
    self.stream.write(data)
    self.stream.flush()
  def writelines(self, datas):
    self.stream.writelines(datas)
    self.stream.flush()
  def __getattr__(self, attr):
    return getattr(self.stream, attr)

sys.stdout = Unbuffered(sys.stdout)

def pad(msg):
  byte = 16 - len(msg) % 16
  return msg + (chr(byte) * byte)

def encrypt(plain, iv, key):
  obj = AES.new(key, AES.MODE_CBC, iv)
  cipher = obj.encrypt(plain)
  return iv + cipher

def decrypt(cipher, key):
  iv = cipher[:16]
  obj = AES.new(key, AES.MODE_CBC, iv)
  plains = obj.decrypt(cipher[16:])
  return plains

def proof_of_work(message, iv, key):
  cips = encrypt(message, iv, key)
  print "ASN : {}".format(b64encode(cips))
  data = raw_input('Student A : ')

  try:
    if decrypt(b64decode(data), key) == "morning_students":
      print "Continue ..."
    else:
      sys.exit(0)
  except:
    print "Fail"
    sys.exit(0)


print """
======================================================
ASN CIPHER SERVICE
======================================================
\n"""

message = "ce_te_ep_joi_nts"
iv      = os.urandom(16)
key     = os.urandom(16)
proof_of_work(message, iv, key)

flag        = open("flag.txt").read()
asn_word    = "Dear my Students, {}".format(os.urandom(32))
num         = 8

candidate = range(256)
random.shuffle(candidate)

garbled_key     = candidate[num:]
garbled_msg     = chr(random.randrange(256)) * (256-num)
garbled_cipher  = "".join(chr(ord(garbled_msg[i]) ^ garbled_key[i]) for i in range(256-num) )

key             = candidate[:num]
key_str         = "".join(map(chr, key)) * 2
enc_asn_word    = encrypt(pad(asn_word), iv, key_str)

print "Garbled Cipher : {}".format(b64encode(garbled_cipher))
print "Encrypted ASN Word : {}".format(b64encode(enc_asn_word))

plain_input = raw_input("Enter Decrypted ASN Word : ")
try:
  if b64decode(plain_input) == asn_word:
    print "Hello Students, Congrats !"
    print "FLAG : {}".format(flag)
  else:
    print "TungTung"
except:
  print "Error -_-"

from Crypto.Cipher import AES
from itertools import permutations
from pwn import *

import base64, os, sys, random

p = process('./asn.py')

def xors(a, b):
    result = ''
    for i, j in zip(a, b):
        result += chr(ord(i) ^ ord(j))
    return result

def unpad(msg):
    byte = ord(msg[-1])
    return msg[:-byte]

def decrypt(cipher, key):
    iv = cipher[:16]
    obj = AES.new(key, AES.MODE_CBC, iv)
    plains = obj.decrypt(cipher[16:])
    return plains

def pow():
    p.recvuntil('ASN : ')
    ASN = base64.b64decode(p.recvline().strip())
    q = xors('morning_students', 'ce_te_ep_joi_nts')
    iv = xors(ASN[:16], q)
    print 'Student A : ' + base64.b64encode(iv + ASN[16:])
    p.sendlineafter('Student A : ', base64.b64encode(iv + ASN[16:]))

pow()

p.recvuntil('Garbled Cipher : ')
garbled_cipher = base64.b64decode(p.recvline().strip())
print 'garbled_cipher : ' + base64.b64encode(garbled_cipher)

p.recvuntil('Encrypted ASN Word : ')
enc_asn_word = base64.b64decode(p.recvline().strip())
print 'enc_asn_word : ' + base64.b64encode(enc_asn_word)

for i in range(256):
    a = xors(garbled_cipher, chr(i)*len(garbled_cipher))
    b = range(256)
    for j in a:
        del b[b.index(ord(j))]
    c = list("".join(map(chr, b)))
    for j in list(permutations(c)):
        key = ''.join(j)*2
        msg = decrypt(enc_asn_word, key)
        if 'Dear my Students, ' in msg:
            print 'Key : ' + repr(key)
            print 'Decrypted ASN Word : ' + base64.b64encode(unpad(msg))
            p.sendlineafter('Enter Decrypted ASN Word : ', base64.b64encode(unpad(msg)))
            p.interactive()
            p.close()
            exit()