run.py

#!/usr/bin/env python3

import os
import sys
import signal
import argparse
import logging
import subprocess
import multiprocessing as mp
import time

import ctypes
import ctypes.util
import os


class CustomFormatter(logging.Formatter):
    grey = "\x1b[38;20m"
    yellow = "\x1b[33;20m"
    green = '\033[92m'
    red = "\x1b[31;20m"
    bold_red = "\x1b[31;1m"
    reset = "\x1b[0m"
    format = "%(asctime)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"

    FORMATS = {
        logging.DEBUG: grey + format + reset,
        logging.INFO: green + format + reset,
        logging.WARNING: yellow + format + reset,
        logging.ERROR: red + format + reset,
        logging.CRITICAL: bold_red + format + reset
    }

    def format(self, record):
        log_fmt = self.FORMATS.get(record.levelno)
        formatter = logging.Formatter(log_fmt)
        return formatter.format(record)

LOG = logging.getLogger("WOFS")
LOG.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
ch.setFormatter(CustomFormatter())
LOG.addHandler(ch)

GRACEFUL_START_SEC = 2

BINS = {
    "wofs_no_crypt": "./bins/{target}/wofs_no_crypt -f {fuse_path} --real {encrypt_path}",
    "wofs_pk":       "./bins/{target}/wofs_pk -f {fuse_path} --init {sk_path} --pk {pk_path} --enc {encrypt_path}",
    "wofs_pk_otp":   "./bins/{target}/wofs_pk_otp -f {fuse_path}  --init {sk_path} --pk {pk_path} --enc {encrypt_path}",
    "wofs_rat":      "./bins/{target}/wofs_rat -f {fuse_path} --init {sk_path} --enc {encrypt_path}",
    "wofs_rat_otp":  "./bins/{target}/wofs_rat_otp -f {fuse_path} --init {sk_path} --enc {encrypt_path}",
    "encfs":         "encfs -f --config encfs_config/encfs.xml --stdinpass {encrypt_path} {fuse_path}",
    "ecryptfs":      None,
    "linux":         None,
}
ECRYPTFS_OPTIONS = [
    "key=passphrase:passphrase_passwd=wofspsswd",
    "ecryptfs_cipher=aes",
    "ecryptfs_key_bytes=16",
    "ecryptfs_unlink_sigs",
    "ecryptfs_enable_filename_crypto=n",
    "ecryptfs_passthrough=n",
    "no_sig_cache",
]
PAYLOAD_SIZES = ("1kB", "1MB", "100MB")

EXPERIMENTS_PATH = "experiments"
EXPERIMENT_ID_TAMPLATE = "{base_path}/{exp_num}"
EXPERIMENT_ID = ""
FUSE_MOUNT_PATH = ""
ENCRYPT_PATH = ""
KEYS_PATH = ""
RESULTS_PATH = ""


libc = ctypes.CDLL(ctypes.util.find_library('c'), use_errno=True)
libc.mount.argtypes = (ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p, ctypes.c_ulong, ctypes.c_char_p)

def mount(source, target, fs, options='', log_file=None):
    subprocess.check_call(
        f'mount -t {fs} {options} {source} {target}',
        shell=True,
        stdout=log_file,
        stderr=log_file
    )
    
def umount(target):
    subprocess.check_call(
        f'umount {target}',
        shell=True,
    )

def parse_params():
    parser = argparse.ArgumentParser(prog='Run', description='Execute WOFS experiments')
    parser.add_argument('-t', '--target', required=True, help="The platform the experiments are executed on. One of [deimos, rpi, notebook]")
    parser.add_argument('-s', '--size', required=True, type=int, help="Number of payload file to be created")
    parser.add_argument('-r', '--runs', required=True, type=int, help="Number of repititions of the experiment")
    
    return parser.parse_args()
            

def log_pid_files(pid, pidstat_log_path, pidio_log_path, pidsmaps_log_path):
    while True:
        try:
            with open(pidstat_log_path, 'a') as pidstat_file:
                with open(os.path.join("/proc", pid, "stat")) as stat_file:
                    pidstat_file.write(f"{time.time_ns()} {stat_file.read()}")
                
            with open(pidio_log_path, 'a') as pidio_file:
                with open(os.path.join("/proc", pid, "io")) as io_file:
                    pidio_file.write(f"{time.time_ns()}\n{io_file.read()}")

            with open(pidsmaps_log_path, 'a') as pidsmaps_file:
                with open(os.path.join("/proc", pid, "smaps")) as smaps_file:
                    pidsmaps_file.write(f"{time.time_ns()}\n{smaps_file.read()}")
        except FileNotFoundError:
            LOG.warning(f"PID files for {pid} is no longer available. Killing pid logging.")
            sys.exit()
            
        time.sleep(0.05)
        

def prepare_dirs():
    LOG.info(f"Creating required folders")
    os.makedirs(FUSE_MOUNT_PATH, exist_ok=True)
    os.makedirs(ENCRYPT_PATH, exist_ok=True)
    os.makedirs(KEYS_PATH, exist_ok=True)
    os.makedirs(RESULTS_PATH, exist_ok=True)
    
    with open(os.path.join(RESULTS_PATH, "clk_tck"), 'w') as clk_tck_file:
        clk_tck_file.write(str(os.sysconf(os.sysconf_names["SC_CLK_TCK"])))
    with open(os.path.join(RESULTS_PATH, "pagesize"), 'w') as pagesize_file:
        pagesize_file.write(str(os.sysconf(os.sysconf_names["SC_PAGESIZE"])))

    mount("none", ENCRYPT_PATH, "nullfs")
      
      
def start_pid_logging(base_log_path, pid, bin):
    LOG.info(f"Starting loggers for {pid}, {bin}")
    stat_log_path = f"{base_log_path}_{bin}.stat"
    io_log_path = f"{base_log_path}_{bin}.io"
    smaps_log_path = f"{base_log_path}_{bin}.smaps"
    pid_process = mp.Process(
        target=log_pid_files,
        args=(str(pid), stat_log_path, io_log_path, smaps_log_path)
    )
    pid_process.start()
    
    return pid_process
                    
    
def start_wofs(target, bin, size, run, log_path):
    KEYS_EXPERIMENT_PATH = os.path.join(KEYS_PATH, f"{bin}_{size}_{run}")
    os.makedirs(KEYS_EXPERIMENT_PATH, exist_ok=True)
    options = {
        "target": target,
        "fuse_path": FUSE_MOUNT_PATH,
        "encrypt_path": ENCRYPT_PATH,
        "sk_path": os.path.join(KEYS_EXPERIMENT_PATH, "sec.key"),
        "pk_path": os.path.join(KEYS_EXPERIMENT_PATH, "pub.key"),
    }
    cmd = BINS[bin].format(**options)
    
    LOG.info(f"Starting {cmd}")
    bin_log_file = open(f"{log_path}.wofs", 'w')
    bin_process = subprocess.Popen(cmd.split(), stdout=bin_log_file, stderr=bin_log_file, stdin=subprocess.PIPE if bin == "encfs" else None)
    if bin == "encfs":
        with open("encfs_config/encfs.passwd", "rb") as encfs_passwd:
            bin_process.stdin.write(encfs_passwd.read().strip())
            bin_process.stdin.close()
    
    return bin_process, bin_log_file


def gen_exp_id():
    global EXPERIMENT_ID, FUSE_MOUNT_PATH, ENCRYPT_PATH, KEYS_PATH, RESULTS_PATH
    
    old_experiments = [int(name) for name in os.listdir(EXPERIMENTS_PATH) if os.path.isdir(f"{EXPERIMENTS_PATH}/{name}") and name.isnumeric()]
    new_id = 1 if len(old_experiments) == 0 else max(old_experiments) + 1
    
    EXPERIMENT_ID = EXPERIMENT_ID_TAMPLATE.format(base_path=EXPERIMENTS_PATH, exp_num=new_id)
    FUSE_MOUNT_PATH = os.path.join(EXPERIMENT_ID, "fuse_mount")
    ENCRYPT_PATH = os.path.join(EXPERIMENT_ID, "encrypt")
    KEYS_PATH = os.path.join(EXPERIMENT_ID, "keys")
    RESULTS_PATH = os.path.join(EXPERIMENT_ID, "results")


if __name__ == '__main__':
    if os.geteuid() != 0:
        LOG.error("You need to have root privileges to run this script.")
        exit(1)
    
    gen_exp_id()
    LOG.info(f"Starting Experiment {EXPERIMENT_ID}")
    
    args = parse_params()
        
    prepare_dirs()
    
    for bin in BINS:
        for size in PAYLOAD_SIZES:
            for run in range(args.runs):
                LOG.info(f"Experiment run: {bin}, {size}, {run}")
                base_log_path = os.path.join(RESULTS_PATH, f'{bin}_{size}_{run}_{args.target}_{args.size}')
                
                if BINS[bin] is not None:
                    bin_process, bin_log_file = start_wofs(args.target, bin, size, run, base_log_path)
                    bin_pid_process = start_pid_logging(base_log_path, bin_process.pid, "bin")
                    
                if bin == "ecryptfs":
                    bin_log_file = open(f"{base_log_path}.wofs", 'w')
                    mount(f"./{ENCRYPT_PATH}", f"./{ENCRYPT_PATH}", "ecryptfs", options='-o ' + ','.join(ECRYPTFS_OPTIONS), log_file=bin_log_file)
                
                time.sleep(GRACEFUL_START_SEC)

                LOG.info("Starting actual copy process")
                dd_log_file = open(f"{base_log_path}.dd", 'w')
                
                if size == "1kB":
                    byte_size = 1000
                elif size == "1MB":
                    byte_size = 1000*1000
                else:
                    byte_size = 1000*1000*100
                
                count = int(args.size / byte_size)
                of_dir = ENCRYPT_PATH if bin in ['linux', "ecryptfs"] else FUSE_MOUNT_PATH
                dd_cmd = ["/usr/bin/dd", "if=/dev/urandom", f"of={of_dir}/{run}", f"bs={size}", f"count={count}"]
                LOG.info(f"Starting dd process: {dd_cmd}")
                dd_process = subprocess.Popen(dd_cmd, stdout=dd_log_file, stderr=dd_log_file)
                
                dd_process.send_signal(signal.SIGSTOP)
                dd_pid_process = start_pid_logging(base_log_path, dd_process.pid, "dd")
                time.sleep(GRACEFUL_START_SEC)
                
                with open(f"{base_log_path}.start", 'w') as start_file:
                    start_file.write(str(time.time_ns()))

                dd_process.send_signal(signal.SIGCONT)
                dd_process.wait()
                
                
                LOG.info("Cleaning up")
                if bin not in ['linux', "ecryptfs"]:
                    bin_process.send_signal(signal.SIGTERM)
                    bin_pid_process.terminate()
                    time.sleep(GRACEFUL_START_SEC)
                    bin_pid_process.close()
                if bin == "ecryptfs":
                    umount(ENCRYPT_PATH)
                if bin != 'linux':
                    bin_log_file.close()
                
                dd_pid_process.terminate()
                time.sleep(GRACEFUL_START_SEC)
                dd_pid_process.close()
                
                for f in os.listdir(ENCRYPT_PATH):
                    os.remove(os.path.join(ENCRYPT_PATH, f))
                for f in os.listdir(FUSE_MOUNT_PATH):
                    os.remove(os.path.join(FUSE_MOUNT_PATH, f))