#!/bin/sh
# Executable patcher
# Copyright (C) 2024-2025  Nguyễn Gia Phong
#
# This file is part of taosc.
#
# Taosc is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Taosc is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with taosc.  If not, see <https://www.gnu.org/licenses/>.

save_exit_code() {
  template="${@:3}"
  cmd="$(printf %s "$template" | sed "s#@@#$2#g")"
  set +e
  AFL_USE_QASAN=1 timeout -k 0 $1 afl-qemu-trace $cmd
  exit_code=$?
  set -e
}

bad() {
  save_exit_code $@ 1>/dev/null 2>&1
  test $exit_code -gt 128 ||
    test $exit_code -ge 124 -a $exit_code -le 127 # timeout
}

if test $# -lt 5
then
  echo Usage: taosc-fix TIMEOUT WORKDIR PROOFS_OF_CONCEPT EXECUTABLE ARG...
  exit 1
fi

timeout=$1
set -eux -o pipefail
wd="$(realpath "$2")"
test -d "$wd"
poc="$(realpath "$3")"
test -d "$poc"
test "$(ls -A "$poc")"
binary="$(realpath "$4")"
test -x "$binary"
bin="$wd/$(basename "$4")"
args="${@:5}"

mkdir -p "$wd"
rm -fr "$wd/poc"
cp -r "$poc" "$wd/poc"
for exploit in "$wd"/poc/*
do
  save_exit_code $timeout "$exploit" "$binary" $args 2>&1 1>/dev/null |
    grep '^    #' |
    grep -F "$binary" |
    sed 's/^    #\([0-9]\+ 0x[0-9a-f]\+\).*$/\1/'
done | sort -n | uniq > "$wd/stack-trace"
test -s "$wd/stack-trace"

(grep '^0 0x[0-9a-f]\+$' "$wd/stack-trace" | sed 's/^0 0x0*//' ||
  true) > "$wd/call-trace"
# Stack trace contains return addresses, not call addresses:
# https://devblogs.microsoft.com/oldnewthing?p=96116
grep -v '^0 0x[0-9a-f]\+$' "$wd/stack-trace" |
  sed 's/^[0-9]\+ 0x0*//' |
  taosc-trace-call "$binary" >> "$wd/call-trace"
test -s "$wd/call-trace"

rm -f "$wd/patch-location"
pushd DATA_DIR 1>/dev/null
trap 'popd 1>/dev/null' EXIT
taosc-scout "$binary" < "$wd/call-trace" |
  while read loc destinations && test ! -f "$wd/patch-location"
  do
    e9tool -100 -M addr=0x$loc -P 'if dest()@jump goto' -o "$bin.jump" "$binary"
    rm -f "$wd/destinations"
    for dest in $destinations
    do
      # In case $wd/poc got poluted
      rm -fr "$wd/poc"
      cp -r "$poc" "$wd/poc"
      for exploit in "$wd/poc"/*
      do
        if TAOSC_DEST=0x$dest bad $timeout "$exploit" "$bin.jump" $args
        then
          continue 2 # next destination
        fi
      done
      echo $loc > "$wd/patch-location"
      echo $dest >> "$wd/destinations"
    done
  done
test -s "$wd/patch-location"
test -s "$wd/destinations"

stack_size=$(taosc-measure-stack "$binary" < "$wd/patch-location")
patch_loc=0x$(< "$wd/patch-location")
e9tool -100 -M addr=$patch_loc -P 'report()@cover' -o "$bin.covered" "$binary"
e9tool -100 -M addr=$patch_loc -P 'log(state)@collect'\
  -o "$bin.collect" "$binary"
e9tool -100 -M addr=$patch_loc -P 'if dest(state)@patch goto'\
  -o "$bin.patched" "$binary"

# TODO: FUZZOLIC's options
fuzzolic -kmprst 90000 -i "$poc" -o "$wd/fuzzolic" -- "$binary" $args ||
  true # FIXME: failing with the same status as the target program
rm -fr "$wd/input"
mkdir -p "$wd/input/benign"
cp -r "$poc" "$wd/input/malicious"
find "$wd/fuzzolic" -name 'test_case_*.dat' -print0 |
  xargs -P$(nproc) -0 -n1 -I @@ \
  taosc-sort-inputs $timeout "$wd"/input/{malicious,benign} @@ \
  "$bin.covered" $args

rm -fr "$wd/values"
mkdir -p "$wd"/values/{benign,bottom,malicious,top}
find "$wd/input" -type f -print0 |
  xargs -I @@ -0 -P$(nproc) -n1 \
  taosc-collect-values $timeout $stack_size "$wd/values" @@ \
  "$bin.collect" $args
find "$wd/values" -type f -print0 | xargs -P$(nproc) -0 -n1 -I @@ mkdir @@.d
find "$wd/values" -maxdepth 2 -type f -print0 |
  xargs -P$(nproc) -0 -n1 -I @@ \
  split -b $((stack_size + 16*8)) @@{,.d/}
find "$wd/values/benign" -mindepth 1 -type d -name '*.d' |
  while read d
  do
    find "$d" -type f -printf '%f\0' |
      xargs -P$(nproc) -0 -n1 -I @@ \
      mv "$d/@@" "$wd/values/bottom/$(basename "$d" .d).@@"
  done
find "$wd/values/malicious" -mindepth 1 -type d -name '*.d' |
  while read d
  do
    find "$d" -type f -printf '%f\0' |
      sort -z |
      head -n -1 -z | # all but last
      xargs -P$(nproc) -0 -n1 -I @@ \
      mv "$d/@@" "$wd/values/bottom/$(basename "$d" .d).@@"
    find "$d" -type f -printf '%f\0' |
      sort -z |
      tail -n 1 -z | # last
      xargs -0 -n1 -I @@ \
      mv "$d/@@" "$wd/values/top/$(basename "$d" .d).@@"
  done
taosc-synth $stack_size "$wd"/values/{bottom,top} > "$wd/predicates"
test -s "$wd/predicates"
# vim: filetype=sh.m4