|
| 1 | +#!/bin/bash |
| 2 | +# Dependency check and install function |
| 3 | +install_missing_dependencies() { |
| 4 | + echo "Checking and installing missing dependencies..." |
| 5 | + local dependencies=("tmux" "python3" "python3-pip") |
| 6 | + local missing_deps=() |
| 7 | + local install_cmd="" |
| 8 | + |
| 9 | + for dep in "${dependencies[@]}"; do |
| 10 | + if ! dpkg -l | grep -qw "$dep"; then |
| 11 | + missing_deps+=("$dep") |
| 12 | + fi |
| 13 | + done |
| 14 | + |
| 15 | + if [ ${#missing_deps[@]} -gt 0 ]; then |
| 16 | + echo "Missing dependencies: ${missing_deps[*]}" |
| 17 | + sudo apt-get update |
| 18 | + |
| 19 | + for dep in "${missing_deps[@]}"; do |
| 20 | + echo "Installing $dep..." |
| 21 | + sudo apt-get install -y $dep |
| 22 | + if [ $? -ne 0 ]; then |
| 23 | + echo "Failed to install $dep. Please try to install it manually." |
| 24 | + exit 1 |
| 25 | + fi |
| 26 | + done |
| 27 | + echo "All dependencies installed successfully." |
| 28 | + else |
| 29 | + echo "All dependencies are already installed." |
| 30 | + fi |
| 31 | +} |
| 32 | + |
| 33 | +# Run the dependency check and installation |
| 34 | +install_missing_dependencies |
| 35 | + |
| 36 | + |
| 37 | +#if [ "$(id -u)" != "0" ]; then echo "This script must be run as root" >&2; exit 1; fi |
| 38 | +# Set SAGE_PATH to the directory of this script |
| 39 | +SAGE_PATH=$(dirname "$(readlink -f "$0")") |
| 40 | + |
| 41 | +# Unlimit the amount of open files |
| 42 | +current_limit=$(ulimit -n) |
| 43 | +[[ "$(lsb_release -rs)" > "22" ]] && export WEBKIT_DISABLE_COMPOSITING_MODE=1 |
| 44 | + |
| 45 | +# Define the desired minimum limit |
| 46 | +desired_limit=80000 |
| 47 | + |
| 48 | +# Path to the GDM3 custom configuration file |
| 49 | +GDM3_CUSTOM_CONF="/etc/gdm3/custom.conf" |
| 50 | +WAYLAND_ENABLED=false |
| 51 | +WATCHDOG_PID=0 |
| 52 | + |
| 53 | +export USE_INVALID_TREE=true |
| 54 | +export PRINT_TIME=true |
| 55 | +export INVALID_TREE_PATH="$SAGE_PATH/invalid_tree/invalid_tree.pickle" |
| 56 | +export RULE_INFO_PATH="$SAGE_PATH/invalid_tree/global_info.pickle" |
| 57 | +export CHROMIUM_PATH="$SAGE_PATH/browser_bins/chrome-asan/chrome" |
| 58 | +export CHROMEDRIVER_PATH="$SAGE_PATH/browser_bins/chromedriver" |
| 59 | +export FIREFOX_PATH="$SAGE_PATH/browser_bins/firefox-asan/firefox" |
| 60 | +export FIREFOXDRIVER_PATH="$SAGE_PATH/browser_bins/firefox-asan/geckodriver" |
| 61 | +export WEBKIT_BINARY_PATH="$SAGE_PATH/browser_bins/webkit/MiniBrowser" |
| 62 | +export WEBKIT_WEBDRIVER_PATH="$SAGE_PATH/browser_bins/webkit/WebKitWebDriver" |
| 63 | +export FREEDOM_PATH="$SAGE_PATH/freedom/main.py" |
| 64 | + |
| 65 | +export FREEDOM_PATH="$SAGE_PATH/freedom/main.py" |
| 66 | +export ORIGAMI_PATH="/home/user/SaGe-Browser-Fuzzer/origami/bin/" |
| 67 | +export FAVOCADO_PATH="/home/user/SaGe-Browser-Fuzzer/favocado/Generator/Run/" |
| 68 | + |
| 69 | +# Check for webkit deps |
| 70 | +dpkg -l | grep -qw libwebkitgtk-6.0-4 || (sudo apt-get update && sudo apt-get install -y libwebkitgtk-6.0-4) |
| 71 | +dpkg -l | grep -qw libavif-dev || (sudo apt-get update && sudo apt-get install -y libavif-dev) |
| 72 | + |
| 73 | +# Function to check Wayland in GDM3 configuration |
| 74 | +check_gdm3_conf() { |
| 75 | + if [ -f "$GDM3_CUSTOM_CONF" ]; then |
| 76 | + if grep -E "^[^#]*WaylandEnable=false" "$GDM3_CUSTOM_CONF" &>/dev/null; then |
| 77 | + echo "Wayland is already disabled in GDM3 configuration." |
| 78 | + elif grep -E "^[^#]*WaylandEnable=true" "$GDM3_CUSTOM_CONF" &>/dev/null; then |
| 79 | + echo "Wayland is enabled in GDM3 configuration. Please disable it to prevent instability when fuzzing." |
| 80 | + WAYLAND_ENABLED=true |
| 81 | + else |
| 82 | + echo "Wayland setting not found in GDM3 configuration. If you are using Wayland, please disable it to prevent instability when fuzzing." |
| 83 | + fi |
| 84 | + else |
| 85 | + echo "GDM3 custom configuration file not found. Skipping..." |
| 86 | + fi |
| 87 | +} |
| 88 | + |
| 89 | +# Function to check Wayland in environment variables |
| 90 | +check_env_vars() { |
| 91 | + if [ "$XDG_SESSION_TYPE" == "wayland" ] || [ "$WAYLAND_DISPLAY" != "" ]; then |
| 92 | + echo "Wayland session detected via environment variables. Please switch to an X11 session to prevent instability when fuzzing." |
| 93 | + WAYLAND_ENABLED=true |
| 94 | + fi |
| 95 | +} |
| 96 | + |
| 97 | +# Execute checks |
| 98 | +check_gdm3_conf |
| 99 | +check_env_vars |
| 100 | + |
| 101 | +# Final decision |
| 102 | +if $WAYLAND_ENABLED; then |
| 103 | + echo "To disable Wayland in GDM3, edit /etc/gdm3/custom.conf and set 'WaylandEnable=false' or comment out the line." |
| 104 | + echo "Then, restart your system or log out and select an X11 session from the login screen." |
| 105 | + exit 1 |
| 106 | +else |
| 107 | + echo "Continuing with the script..." |
| 108 | +fi |
| 109 | + |
| 110 | +# Check if apport is installed |
| 111 | +if dpkg-query -W -f='${Status}' apport 2>/dev/null | grep -q "install ok installed"; then |
| 112 | + echo "Apport is currently installed on your system. This can get messy." |
| 113 | + echo "Please uninstall apport before proceeding with this script." |
| 114 | + echo "You can uninstall apport by running: sudo apt-get remove --purge apport" |
| 115 | + exit 1 |
| 116 | +else |
| 117 | + echo "apport is not installed, proceeding..." |
| 118 | +fi |
| 119 | + |
| 120 | +# Function to kill all spawned processes, browser processes, and any process from SAGE_PATH |
| 121 | +cleanup() { |
| 122 | + echo "Terminating all spawned processes, browser processes, any process from SAGE_PATH, and the watchdog..." |
| 123 | + |
| 124 | + # Kills jobs spawned by this script |
| 125 | + kill $(jobs -p) 2>/dev/null |
| 126 | + |
| 127 | + # Kills all child processes spawned by this script |
| 128 | + pkill -P $$ 2>/dev/null |
| 129 | + |
| 130 | + # Explicitly kill processes started from SAGE_PATH |
| 131 | + pkill -f "$SAGE_PATH" 2>/dev/null && pkill tmux |
| 132 | + |
| 133 | + # If watchdog is running, kill it |
| 134 | + if [ $WATCHDOG_PID -ne 0 ]; then |
| 135 | + kill -9 $WATCHDOG_PID 2>/dev/null |
| 136 | + echo "Watchdog (PID $WATCHDOG_PID) terminated." |
| 137 | + fi |
| 138 | +} |
| 139 | + |
| 140 | +# Function to kill old processes from SAGE_PATH before starting new ones |
| 141 | +kill_old_processes() { |
| 142 | + echo "Killing old processes started from $SAGE_PATH..." |
| 143 | + pkill -f "$SAGE_PATH" 2>/dev/null |
| 144 | +} |
| 145 | + |
| 146 | +# Function to monitor system memory, restart browser bins/drivers if needed, and auto-kill processes after a set timeout |
| 147 | +watchdog() { |
| 148 | + local start_time=$(date +%s) |
| 149 | + local browser_bin_names=("chrome" "chromedriver" "firefox" "geckodriver" "MiniBrowser" "WebKitWebDriver") |
| 150 | + local exclude_utilities=("tmux" "tree" "watch" "lolcat" "stat" "tail" "find" "comm" "basename" "btop" "ifne" "grep" "ps") |
| 151 | + |
| 152 | + while :; do |
| 153 | + local current_time=$(date +%s) |
| 154 | + local elapsed_time=$((current_time - start_time)) |
| 155 | + |
| 156 | + if [[ -n "$TIMER_PURGE" && "$elapsed_time" -ge "$TIMER_PURGE" ]]; then |
| 157 | + echo "Timer purge limit reached. Performing cleanup..." |
| 158 | + for bin_name in "${browser_bin_names[@]}"; do |
| 159 | + for pid in $(pgrep -f "$bin_name"); do |
| 160 | + local cmd=$(ps -p $pid -o comm=) |
| 161 | + local exclude=false |
| 162 | + for exclude_cmd in "${exclude_utilities[@]}"; do |
| 163 | + if [[ "$cmd" == *"$exclude_cmd"* ]]; then |
| 164 | + exclude=true |
| 165 | + break |
| 166 | + fi |
| 167 | + done |
| 168 | + if [ "$exclude" = false ]; then |
| 169 | + kill -9 $pid 2>/dev/null |
| 170 | + fi |
| 171 | + done |
| 172 | + done |
| 173 | + start_time=$(date +%s) |
| 174 | + fi |
| 175 | + |
| 176 | + local free_ram=$(awk '/MemAvailable/ {print $2}' /proc/meminfo) |
| 177 | + local available_ram_mb=$((free_ram/1024)) |
| 178 | + if [[ "$available_ram_mb" -lt 2000 ]]; then |
| 179 | + for bin_name in "${browser_bin_names[@]}"; do |
| 180 | + for pid in $(pgrep -f "$bin_name"); do |
| 181 | + local cmd=$(ps -p $pid -o comm=) |
| 182 | + local exclude=false |
| 183 | + for exclude_cmd in "${exclude_utilities[@]}"; do |
| 184 | + if [[ "$cmd" == *"$exclude_cmd"* ]]; then |
| 185 | + exclude=true |
| 186 | + break |
| 187 | + fi |
| 188 | + done |
| 189 | + if [ "$exclude" = false ]; then |
| 190 | + kill -9 $pid 2>/dev/null |
| 191 | + fi |
| 192 | + done |
| 193 | + done |
| 194 | + fi |
| 195 | + sleep 5 |
| 196 | + done |
| 197 | +} |
| 198 | + |
| 199 | +# Handle Ctrl-C (SIGINT) and script exit (EXIT) |
| 200 | +trap cleanup SIGINT EXIT |
| 201 | + |
| 202 | +# Initialize the fuzzer variable with a default value |
| 203 | +FUZZER="sage" |
| 204 | + |
| 205 | +# Initialize browsers and their instance counts |
| 206 | +declare -A BROWSER_INSTANCES |
| 207 | +KILL_OLD=false |
| 208 | +WATCHDOG_ENABLED=false |
| 209 | +TIMER_PURGE="" |
| 210 | + |
| 211 | +# General output directory for logs |
| 212 | +GENERAL_OUTPUT_DIR=$SAGE_PATH/output |
| 213 | +mkdir -p "$GENERAL_OUTPUT_DIR" |
| 214 | +LOG_FILE="$GENERAL_OUTPUT_DIR/main.log" |
| 215 | + |
| 216 | +# Function to check and trim the log file |
| 217 | +# Set default log size limit (in MB) |
| 218 | +LOG_SIZE_LIMIT=${LOG_SIZE_LIMIT:-500} |
| 219 | +USE_LOLCAT=true |
| 220 | +trim_log_file() { |
| 221 | + local max_size=$((LOG_SIZE_LIMIT * 1024 * 1024)) # Convert MB to bytes |
| 222 | + while true; do |
| 223 | + local file_size=$(stat -c%s "$LOG_FILE" 2>/dev/null || echo 0) |
| 224 | + if (( file_size > max_size )); then |
| 225 | + echo "Trimming $LOG_FILE (size: $file_size, limit: $max_size)" |
| 226 | + tail -c "$max_size" "$LOG_FILE" > "$LOG_FILE.tmp" |
| 227 | + if [ -s "$LOG_FILE.tmp" ]; then |
| 228 | + mv "$LOG_FILE.tmp" "$LOG_FILE" |
| 229 | + else |
| 230 | + echo "Temporary file not created or is empty, skipping move operation." |
| 231 | + fi |
| 232 | + fi |
| 233 | + sleep 60 # Check every 60 seconds |
| 234 | + done |
| 235 | +} |
| 236 | + |
| 237 | +# Start trimming log file in the background |
| 238 | +trim_log_file & |
| 239 | + |
| 240 | +# Check command line arguments |
| 241 | +while (( "$#" )); do |
| 242 | + case "$1" in |
| 243 | + --firefox|--webkitgtk|--chromium) |
| 244 | + BROWSER=${1#--} |
| 245 | + if [[ -n "$2" && "$2" =~ ^[0-9]+$ ]]; then |
| 246 | + BROWSER_INSTANCES[$BROWSER]=$2 |
| 247 | + shift 2 |
| 248 | + else |
| 249 | + echo "Error: Expected a number of instances after $1" |
| 250 | + exit 1 |
| 251 | + fi |
| 252 | + ;; |
| 253 | + --fuzzer) |
| 254 | + if [[ -n "$2" && "$2" =~ ^(domato|minerva|freedom|sage|favocado)$ ]]; then |
| 255 | + FUZZER=$2 |
| 256 | + shift 2 |
| 257 | + else |
| 258 | + echo "Error: Unsupported fuzzer. Supported fuzzers are domato, minerva, freedom, sage, and favocado." |
| 259 | + exit 1 |
| 260 | + fi |
| 261 | + ;; |
| 262 | + --kill-old) |
| 263 | + KILL_OLD=true |
| 264 | + shift |
| 265 | + ;; |
| 266 | + --watchdog) |
| 267 | + WATCHDOG_ENABLED=true |
| 268 | + shift |
| 269 | + ;; |
| 270 | + --timerpurge) |
| 271 | + if [[ -n "$2" && "$2" =~ ^[0-9]+$ ]]; then |
| 272 | + TIMER_PURGE=$2 |
| 273 | + shift 2 |
| 274 | + else |
| 275 | + echo "Error: Expected a numeric value for --timerpurge" |
| 276 | + exit 1 |
| 277 | + fi |
| 278 | + ;; |
| 279 | + *) |
| 280 | + echo "Unsupported option: $1" |
| 281 | + exit 1 |
| 282 | + ;; |
| 283 | + esac |
| 284 | +done |
| 285 | + |
| 286 | +# If --kill-old was specified, kill old processes |
| 287 | +if [ "$KILL_OLD" = true ]; then |
| 288 | + kill_old_processes |
| 289 | +fi |
| 290 | + |
| 291 | +# If --watchdog was specified, start the watchdog function in the background |
| 292 | +if [ "$WATCHDOG_ENABLED" = true ]; then |
| 293 | + watchdog & |
| 294 | + WATCHDOG_PID=$! |
| 295 | +fi |
| 296 | + |
| 297 | +# Function to generate a unique output directory |
| 298 | +generate_unique_output_dir() { |
| 299 | + local browser_name=$1 |
| 300 | + local datetime=$(date +%Y-%m-%d-%H-%M-%S) |
| 301 | + local uid=$(uuidgen | cut -d'-' -f1) # Generates a short UID from uuidgen |
| 302 | + local output_dir="$SAGE_PATH/output/$browser_name/$datetime-$uid" |
| 303 | + echo $output_dir |
| 304 | +} |
| 305 | + |
| 306 | +# Replace the section where PYTHON_OUTPUT_DIR is set with the below code |
| 307 | +# This ensures a unique directory is created for each session |
| 308 | + |
| 309 | +# Initialize browsers and their instance counts |
| 310 | +declare -A BROWSER_INSTANCES |
| 311 | +KILL_OLD=false |
| 312 | +WATCHDOG_ENABLED=false |
| 313 | +TIMER_PURGE="" |
| 314 | + |
| 315 | +# General output directory for logs and modification for unique output directories |
| 316 | +GENERAL_OUTPUT_DIR=$SAGE_PATH/output |
| 317 | +mkdir -p "$GENERAL_OUTPUT_DIR" |
| 318 | +LOG_FILE="$GENERAL_OUTPUT_DIR/main.log" |
| 319 | + |
| 320 | +# Start trimming log file in the background function (if previously defined in your script) |
| 321 | +trim_log_file & |
| 322 | + |
| 323 | +# Check command line arguments and setup (Keep your existing argument parsing logic here) |
| 324 | + |
| 325 | +# Logic to start fuzzing sessions with unique output directories |
| 326 | +for BROWSER in "${!BROWSER_INSTANCES[@]}" |
| 327 | +do |
| 328 | + NUM_INSTANCES=${BROWSER_INSTANCES[$BROWSER]} |
| 329 | + # Generate a unique output directory for this session |
| 330 | + PYTHON_OUTPUT_DIR=$(generate_unique_output_dir $BROWSER) |
| 331 | + mkdir -p "$PYTHON_OUTPUT_DIR" |
| 332 | + |
| 333 | + # Start main.py with specified parameters and redirect output to both the log file and terminal |
| 334 | + python3 $SAGE_PATH/main.py -t 50000 -b $BROWSER -p $NUM_INSTANCES --fuzzer $FUZZER -o $PYTHON_OUTPUT_DIR 2>&1 | tee -a "$LOG_FILE" & |
| 335 | + |
| 336 | + # Record the start time and write it to a file |
| 337 | + echo $(date +%s) > "$PYTHON_OUTPUT_DIR/start_time.txt" |
| 338 | +done |
| 339 | + |
| 340 | +# Wait for all background processes to finish |
| 341 | +wait |
0 commit comments