研究室内CTFで出題したCrypto問題のサービス

service.py

#!/usr/bin/python
import hashlib
import base64
import re
from secret import FLAG, SALT

def construct(data):
    ret = {}
    regs = data.split('|')
    for reg in regs:
        key, value = reg.split(':')
        ret[key] = value
    return ret

def sign(user):
    data = 'priv:student|user:' + user
    cert = data + hashlib.md5(SALT + data).hexdigest()
    return base64.b64encode(cert)

def login(cert):
    try:
        data = base64.b64decode(cert)
    except TypeError:
        return False
    hash = data[-32:]
    data = data[:-32]
    if hashlib.md5(SALT + data).hexdigest() != hash:
        return False
    try:
        ret = construct(data)
    except ValueError:
        return False
    return ret

def action(data):
    print("[+] Welcome, " + data['user'] + ".")
    while True:
        print("[1]Info / [2]FLAG / [3]Exit")
        try:
            n = int(raw_input(">> "))
        except ValueError:
            continue
        if n == 1:
            print("Username : " + data['user'])
            print("Privilege: " + data['priv'])
        elif n == 2:
            if data['priv'] == 'teacher':
                print("[+] The flag is " + FLAG)
            else:
                print("[-] This action can be issued only by teachers.")
        else:
            print("See you, " + user + ".")
            break
    return True

if __name__ == '__main__':
    print("+------------------------------+")
    print("| Ocamlab Digital Sign Service |")
    print("+------------------------------+")
    print("- Students must get certification and use it to login.")
    print("- Teachers can use pre-distributed certification.\n")
    while True:
        print("[1]Sign / [2]Login")
        try:
            n = int(raw_input(">> "))
        except ValueError:
            continue
        if n != 1 and n != 2: continue
        if n == 1:
            while True:
                user = raw_input("Username: ")
                if re.match('^[\w]+$', user):
                    break
                print("[-] Invalid username")
            print("[+] Hello, " + user + ".")
            cert = sign(user)
            print("Your Certification: " + cert)
        else:
            cert = raw_input("Certification: ")
            data = login(cert)
            if data:
                action(data)
                break
            else:
                print("[-] Invalid certification")