verify_exploit.py

#!/usr/bin/env python
###############################################################################
# Git-based CTF
###############################################################################
#
# Author: SeongIl Wi <seongil.wi@kaist.ac.kr>
#         Jaeseung Choi <jschoi17@kaist.ac.kr>
#         Sang Kil Cha <sangkilc@kaist.ac.kr>
#
# Copyright (c) 2018 SoftSec Lab. KAIST
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import sys
import os
import json
from utils import random_string, docker_cleanup, base_dir, load_config
from utils import prompt_checkout_warning, get_dirname, print_and_log
from cmd import run_command
from git import checkout
from crypto import encrypt_exploit

#-*- coding: utf-8 -*-

# TODO : Get these values via cmdline option
SERVICE_IP = "127.0.0.1"
SERVICE_PORT = 4000

def start_service(service_dir, branch, container_name, flag_str, log=None):

    log= print_and_log("[*] Starting service from %s (branch '%s')" % \
            (service_dir, branch), log)

    checkout(service_dir, branch)

    # Update flag file
    flag_path = os.path.join(service_dir, "flag") # Assumption in template
    if not os.path.isfile(flag_path):
        log = print_and_log("[*] 'flag' file not found in %s" % service_dir, log)
        return False, log
    with open(flag_path, "w") as flag_file:
        flag_file.write(flag_str)

    # Run the service
    script = os.path.join(base_dir(), "setup_service.sh")
    cmdline = \
        "%s %s %d %d" % (script, container_name, SERVICE_PORT, SERVICE_PORT)
    output, err, e = run_command(cmdline, service_dir)
    if e != 0:
        log = print_and_log("[*] Failed to start service", log)
        log = print_and_log(err, log)
        log = print_and_log("==========================", log)
        return False, log
    if log is not None:
        log = log + output

    log = print_and_log("[*] Started service successfully", log)
    return True, log

def run_exploit(exploit_dir, container_name, timeout, log=None):
    log = print_and_log("[*] Running exploit", log)

    script = os.path.join(base_dir(), "launch_exploit.sh")
    cmdline = \
      "%s %s %s %d %d" % \
      (script, container_name, SERVICE_IP, SERVICE_PORT, timeout)
    output, err, e = run_command(cmdline, exploit_dir)
    if log is not None:
        try:
            log = log.decode('utf_8', 'ignore')
            output = output.decode('utf-8', 'ignore')
            log = log + output
        except UnicodeDecodeError:
            pass

    if e != 0:
        log = print_and_log("[*] Failed to run exploit", log)
        log = print_and_log(err, log)
        log = print_and_log("==========================", log)
        return None, log

    # Exploit prints out the flag string at the end.
    tokens = filter(None, output.split('\n')) # Filter out empty strings
    flag_candidate = filter(None, tokens)[-1] # Read the last line
    return flag_candidate, log

def verify_exploit(exploit_dir, service_dir, branch, timeout, config,
                   encrypt=False, log=None):
    if not os.path.isdir(exploit_dir) :
        print "[*] Exploit directory '%s' does not exist" % exploit_dir
        return False, log

    if not os.path.isdir(service_dir) :
        print "[*] Service directory '%s' does not exist" % service_dir
        return False, log

    # Create random flag value
    flag = random_string(10)

    # Start the service
    service_dirname = get_dirname(service_dir)
    service_container_name = "%s-%s" % (service_dirname, branch)
    result, log = start_service(service_dir, branch, service_container_name, \
            flag, log=log)
    if not result:
        return False, log

    # Run the exploit
    exploit_dirname = get_dirname(exploit_dir)
    exploit_container_name = "exploit-%s" % branch
    exploit_result, log = run_exploit(exploit_dir, exploit_container_name, \
            timeout, log=log)

    # Clean up containers
    docker_cleanup(service_container_name)
    docker_cleanup(exploit_container_name)

    log = print_and_log("[*] Exploit returned : %s" % exploit_result, log)
    log = print_and_log("[*] Solution flag : %s" % flag, log)
    if exploit_result == flag:
        print "[*] Exploit worked successfully"
        if encrypt:
            print "[*] Encrypting the verified exploit"
            # Set your own team as target team, and signer is not needed.
            target_team = config["player_team"]
            encrypted_file = encrypt_exploit(exploit_dir, target_team, config)
            if encrypted_file is None:
                print "[*] Failed to encrypt exploit"
            else:
                print "[*] Your exploit is encrypted in %s" % encrypted_file
                print "[*] Now you may commit and push this encrypted exploit "\
                      "to the corresponding branch of your service repository"
        return True, log
    else:
        log = print_and_log("[*] Exploit returned a wrong flag string", log)
        return False, log

if __name__ == "__main__":
    if len(sys.argv) != 6:
        print ("Usage: %s " % sys.argv[0] +
               "[exploit dir] [service dir] [branch] [timeout] [config]")
        sys.exit()

    exploit_dir = sys.argv[1]
    service_dir = sys.argv[2]
    branch = sys.argv[3]
    timeout = int(sys.argv[4])
    config_file = sys.argv[5]
    prompt_checkout_warning(service_dir)
    config = load_config(config_file)
    verify_exploit(exploit_dir, service_dir, branch, timeout, config)