#!/usr/bin/env bash

# Usage: run option -h to get help

# BT auto detection
#   Shall we look for white HC-06-USB dongle ?
FINDBTDONGLE=true
#   Shall we look for rfcomm interface ?
FINDBTRFCOMM=true
#   Shall we look for registered BT device ? (Linux only)
FINDBTDIRECT=true

PM3PATH=$(dirname "$0")
EVALENV=""
FULLIMAGE="fullimage.elf"
BOOTIMAGE="bootrom.elf"

#Skip check if --list is used
if [ ! "$1" == "--list" ]; then
    # try pm3 dirs in current repo workdir
    if [ -d "$PM3PATH/client/" ]; then
        if [ -x "$PM3PATH/client/proxmark3" ]; then
            CLIENT="$PM3PATH/client/proxmark3"
        elif [ -x "$PM3PATH/client/build/proxmark3" ]; then
            CLIENT="$PM3PATH/client/build/proxmark3"
        else
            echo >&2 "[!!] In devel workdir but no executable found, did you compile it?"
            exit 1
        fi
        # Devel mode: point to workdir pm3.py module
        EVALENV+=" PYTHONPATH=$PM3PATH/client/src"
    # try install dir
    elif [ -x "$PM3PATH/proxmark3" ]; then
        CLIENT="$PM3PATH/proxmark3"
        EVALENV+=" PYTHONPATH=$PM3PATH/../share/proxmark3/pyscripts/"
        #  or /usr/[local/]lib/python3/dist-packages/pm3.py ?
    else
    # hope it's installed somehow, still not sure where fw images and pm3.py are...
        CLIENT="proxmark3"
    fi
fi

# LeakSanitizer suppressions
if [ -e .lsan_suppressions ]; then
  EVALENV+=" LSAN_OPTIONS=suppressions=.lsan_suppressions"
fi
if [ "$EVALENV" != "" ]; then
  EVALENV="export $EVALENV"
fi
PM3LIST=()
SHOWLIST=false

function get_pm3_list_Linux {
    N=$1
    PM3LIST=()
    if [ ! -c "/dev/tty0" ]; then
        echo >&2 "[!!] Script cannot access /dev/ttyXXX files, insufficient privileges"
        exit 1
    fi
    for DEV in $(find /dev/ttyACM* 2>/dev/null); do
        if command -v udevadm >/dev/null; then
            if udevadm info -q property -n "$DEV" | grep -q "ID_VENDOR=proxmark.org"; then
                PM3LIST+=("$DEV")
                if [ ${#PM3LIST[*]} -ge "$N" ]; then
                    return
                fi
            fi
        else
            if grep -q "proxmark.org" "/sys/class/tty/${DEV#/dev/}/../../../manufacturer" 2>/dev/null; then
                PM3LIST+=("$DEV")
                if [ ${#PM3LIST[*]} -ge "$N" ]; then
                    return
                fi
            fi
        fi
    done
    if $FINDBTDONGLE; then
        # check if the HC-06-USB white dongle is present (still, that doesn't tell us if it's paired with a Proxmark3...)
        for DEV in $(find /dev/ttyUSB* 2>/dev/null); do
            if command -v udevadm >/dev/null; then
                if udevadm info -q property -n "$DEV" | grep -q "ID_MODEL=CP2104_USB_to_UART_Bridge_Controller"; then
                    PM3LIST+=("$DEV")
                    if [ ${#PM3LIST[*]} -ge "$N" ]; then
                        return
                    fi
                fi
            else
                if grep -q "DRIVER=cp210x" "/sys/class/tty/${DEV#/dev/}/../../uevent" 2>/dev/null; then
                    PM3LIST+=("$DEV")
                    if [ ${#PM3LIST[*]} -ge "$N" ]; then
                        return
                    fi
                fi
            fi
        done
    fi
    if $FINDBTRFCOMM; then
        # check if the MAC of a Proxmark3 was bound to a local rfcomm interface
        # (on OSes without deprecated rfcomm and hcitool, the loop will be simply skipped)
        for DEVMAC in $(rfcomm -a 2>/dev/null | grep " 20:19:0[45]" | sed 's/^\(.*\): \([0-9:]*\) .*/\1@\2/'); do
            DEV=${DEVMAC/@*/}
            MAC=${DEVMAC/*@/}
            # check which are Proxmark3 and, side-effect, if they're actually present
            if hcitool name "$MAC" | grep -q "PM3"; then
                PM3LIST+=("/dev/$DEV")
                if [ ${#PM3LIST[*]} -ge "$N" ]; then
                    return
                fi
            fi
        done
    fi
    if $FINDBTDIRECT; then
        # check if the MAC of a Proxmark3 was registered in the known devices
        for MAC in $(dbus-send --system --print-reply --type=method_call --dest='org.bluez' '/' org.freedesktop.DBus.ObjectManager.GetManagedObjects 2>/dev/null|\
                         awk '/"Address"/{getline;gsub(/"/,"",$3);a=$3}/Name/{getline;if (/PM3_RDV4/) print a}'); do
            PM3LIST+=("bt:$MAC")
        done
        # we don't probe the device so there is no guarantee the device is actually present
    fi
}

function get_pm3_list_macOS {
    N=$1
    PM3LIST=()
    for DEV in $(ioreg -r -c "IOUSBHostDevice" -l | awk -F '"' '
        $2=="USB Vendor Name"{b=($4=="proxmark.org")}
        b==1 && $2=="IODialinDevice"{print $4}'); do
        PM3LIST+=("$DEV")
        if [ ${#PM3LIST[*]} -ge "$N" ]; then
            return
        fi
    done
}

function get_pm3_list_Windows {
    N=$1
    PM3LIST=()
    # Need to look for this first,  the call to Win32_serialport "crashes" then native bt serial port.  Don't ask why.
    #BT direct SERIAL PORTS (COM)
    if $FINDBTRFCOMM; then
        for DEV in $(wmic /locale:ms_409 path Win32_PnPEntity Where "Caption LIKE '%Bluetooth%(COM%'" Get Name 2> /dev/null | awk -b 'match($0,/(COM[0-9]+)/,m){print m[1]}'); do
            DEV=${DEV/ */}
            PM3LIST+=("$DEV")
            if [ ${#PM3LIST[*]} -ge "$N" ]; then
                return
            fi
        done
    fi

    # Normal SERIAL PORTS (COM)
    for DEV in $(wmic /locale:ms_409 path Win32_SerialPort Where "PNPDeviceID LIKE '%VID_9AC4&PID_4B8F%' Or PNPDeviceID LIKE '%VID_2D2D&PID_504D%'" Get DeviceID 2>/dev/null | awk -b '/^COM/{print $1}'); do
        DEV=${DEV/ */}
        #prevent soft bricking when using pm3-flash-all on an outdated bootloader
        if [ $(basename -- "$0") = "pm3-flash-all" ]; then
            if [ ! $(wmic /locale:ms_409 path Win32_SerialPort Where "DeviceID='$DEV'" Get PNPDeviceID 2>/dev/null | awk -b '/^USB/{print $1}') = "USB\VID_9AC4&PID_4B8F\ICEMAN" ]; then
                echo -e "\033[0;31m[!] Using pm3-flash-all on an oudated bootloader, use pm3-flash-bootrom first!"
                exit 1
            fi
        fi
        #Prioritise USB connections
        PM3LIST=("$DEV" "${PM3LIST[@]}")
        if [ ${#PM3LIST[*]} -ge "$N" ]; then
            return
        fi
    done

    #white BT dongle SERIAL PORTS (COM)
    if $FINDBTDONGLE; then
        for DEV in $(wmic /locale:ms_409 path Win32_SerialPort Where "PNPDeviceID LIKE '%VID_10C4&PID_EA60%'" Get DeviceID 2>/dev/null | awk -b '/^COM/{print $1}'); do
            DEV=${DEV/ */}
            PM3LIST+=("$DEV")
            if [ ${#PM3LIST[*]} -ge "$N" ]; then
                return
            fi
        done
    fi
}

function get_pm3_list_WSL {
    N=$1
    PM3LIST=()

    # Need to look for this first,  the call to Win32_serialport "crashes" then native bt serial port.  Don't ask why.
    #BT direct SERIAL PORTS (COM)
    if $FINDBTRFCOMM; then
        for DEV in $($PSHEXE -command "Get-CimInstance -ClassName Win32_PnPEntity | Where-Object Caption -like 'Standard Serial over Bluetooth link (COM*' | Select Name" 2> /dev/null | sed -nr 's#.*\bCOM([0-9]+)\b.*#/dev/ttyS\1#p'); do
            # ttyS counterpart takes some more time to appear
            if [ -e "$DEV" ]; then
                PM3LIST+=("$DEV")
                if [ ! -w "$DEV" ]; then
                    echo "[!] Let's give users read/write access to $DEV"
                    sudo chmod 666 "$DEV"
                fi
                if [ ${#PM3LIST[*]} -ge "$N" ]; then
                    return
                fi
            fi

        done
    fi

    # Normal SERIAL PORTS (COM)
    for DEV in $($PSHEXE -command "Get-CimInstance -ClassName Win32_serialport | Where-Object {\$_.PNPDeviceID -like '*VID_9AC4&PID_4B8F*' -or \$_.PNPDeviceID -like '*VID_2D2D&PID_504D*'} | Select -expandproperty DeviceID" 2>/dev/null | tr -dc '[:print:]'); do
        _comport=$DEV
        DEV=$(echo $DEV | sed -nr 's#^COM([0-9]+)\b#/dev/ttyS\1#p')
        # ttyS counterpart takes some more time to appear
        if [ -e "$DEV" ]; then
            #prevent soft bricking when using pm3-flash-all on an outdated bootloader
            if [ $(basename -- "$0") = "pm3-flash-all" ]; then
                if [ ! $($PSHEXE -command "Get-CimInstance -ClassName Win32_serialport | Where-Object {\$_.DeviceID -eq '$_comport'} | Select -expandproperty PNPDeviceID" 2>/dev/null | tr -dc '[:print:]') = "USB\VID_9AC4&PID_4B8F\ICEMAN" ]; then
                    echo -e "\033[0;31m[!] Using pm3-flash-all on an oudated bootloader, use pm3-flash-bootrom first!"
                    exit 1
                fi
            fi
            #Prioritise USB connections
            PM3LIST=("$DEV" "${PM3LIST[@]}")
            if [ ! -w "$DEV" ]; then
                echo "[!] Let's give users read/write access to $DEV"
                sudo chmod 666 "$DEV"
            fi
            if [ ${#PM3LIST[*]} -ge "$N" ]; then
                return
            fi
        fi
    done

    #white BT dongle SERIAL PORTS (COM)
    if $FINDBTDONGLE; then
        for DEV in $($PSHEXE -command "Get-CimInstance -ClassName Win32_serialport | Where-Object PNPDeviceID -like '*VID_10C4&PID_EA60*' | Select DeviceID" 2>/dev/null | sed -nr 's#^COM([0-9]+)\b#/dev/ttyS\1#p'); do
            # ttyS counterpart takes some more time to appear
            if [ -e "$DEV" ]; then
                PM3LIST+=("$DEV")
                if [ ! -w "$DEV" ]; then
                    echo "[!] Let's give users read/write access to $DEV"
                    sudo chmod 666 "$DEV"
                fi
                if [ ${#PM3LIST[*]} -ge "$N" ]; then
                    return
                fi
            fi

        done
    fi
}

SCRIPT=$(basename -- "$0")

if [ "$SCRIPT" = "pm3" ]; then
  CMD() { eval "$EVALENV"; $CLIENT "$@"; }
  HELP() {
      cat << EOF

Quick helper script for proxmark3 client when working with a Proxmark3 device

Description:
    The usage is the same as for the proxmark3 client, with the following differences:
     * the correct port name will be automatically guessed;
     * the script will wait for a Proxmark to be connected (same as option -w of the client).
    You can also specify a first option -n N to access the Nth Proxmark3 connected.
    To see a list of available ports, use --list.

Usage:
    $SCRIPT [-n <N>] [<any other proxmark3 client option>]
    $SCRIPT [--list] [-h|--help] [-hh|--helpclient]
    $SCRIPT [-o|--offline]


Arguments:
    -h/--help        this help
    -hh/--helpclient proxmark3 client help (the script will forward these options)
    --list           list all detected com ports
    -n <N>           connect device referred to the N:th number on the --list output
    -o/--offline     shortcut to use directly the proxmark3 client without guessing ports

Samples:
    ./$SCRIPT                       -- Auto detect/ select com port in the following order BT, USB/CDC, BT DONGLE
    ./$SCRIPT -p /dev/ttyACM0       -- connect to port /dev/ttyACM0
    ./$SCRIPT -n 2                  -- use second item from the --list output
    ./$SCRIPT -c 'lf search' -i     -- run command and stay in client once completed


EOF
  }
elif [ "$SCRIPT" = "pm3-flash" ]; then
  FINDBTDONGLE=false
  FINDBTRFCOMM=false
  FINDBTDIRECT=false
  CMD() {
      ARGS=("--port" "$1" "--flash")
      shift;
      while [ "$1" != "" ]; do
          if [ "$1" == "-b" ]; then
              ARGS+=("--unlock-bootloader")
          else
              ARGS+=("--image" "$1")
          fi
          shift;
      done
      $CLIENT "${ARGS[@]}";
  }
  HELP() {
      cat << EOF
Quick helper script for flashing a Proxmark device via USB

Description:
    The usage is similar to the old proxmark3-flasher binary, except that the correct port name will be automatically guessed.
    You can also specify a first option -n N to access the Nth Proxmark3 connected on USB.
    If this doesn't work, you'll have to use manually the proxmark3 client, see "$CLIENT -h".
    To see a list of available ports, use --list.

Usage:
    $SCRIPT [-n <N>] [-b] image.elf [image.elf...]
    $SCRIPT --list

Options:
    -b         Enable flashing of bootloader area (DANGEROUS)

Example:
     $SCRIPT -b bootrom.elf fullimage.elf
EOF
  }
elif [ "$SCRIPT" = "pm3-flash-all" ]; then
  FINDBTDONGLE=false
  FINDBTRFCOMM=false
  FINDBTDIRECT=false
  CMD() { $CLIENT "--port" "$1" "--flash" "--unlock-bootloader" "--image" "$BOOTIMAGE" "--image" "$FULLIMAGE"; }
  HELP() {
      cat << EOF
Quick helper script for flashing a Proxmark device via USB

Description:
    The correct port name will be automatically guessed and the stock bootloader and firmware image will be flashed.
    You can also specify a first option -n N to access the Nth Proxmark3 connected on USB.
    If this doesn't work, you'll have to use manually the proxmark3 client, see "$CLIENT -h".
    To see a list of available ports, use --list.

Usage:
    $SCRIPT [-n <N>]
    $SCRIPT --list
EOF
  }
elif [ "$SCRIPT" = "pm3-flash-fullimage" ]; then
  FINDBTDONGLE=false
  FINDBTRFCOMM=false
  FINDBTDIRECT=false
  CMD() { $CLIENT "--port" "$1" "--flash" "--image" "$FULLIMAGE"; }
  HELP() {
      cat << EOF
Quick helper script for flashing a Proxmark device via USB

Description:
    The correct port name will be automatically guessed and the stock firmware image will be flashed.
    You can also specify a first option -n N to access the Nth Proxmark3 connected on USB.
    If this doesn't work, you'll have to use manually the proxmark3 client, see "$CLIENT -h".
    To see a list of available ports, use --list.

Usage:
    $SCRIPT [-n <N>]
    $SCRIPT --list
EOF
  }
elif [ "$SCRIPT" = "pm3-flash-bootrom" ]; then
  FINDBTDONGLE=false
  FINDBTRFCOMM=false
  FINDBTDIRECT=false
  CMD() { $CLIENT "--port" "$1" "--flash" "--unlock-bootloader" "--image" "$BOOTIMAGE"; }
  HELP() {
      cat << EOF
Quick helper script for flashing a Proxmark device via USB

Description:
    The correct port name will be automatically guessed and the stock bootloader will be flashed.
    You can also specify a first option -n N to access the Nth Proxmark3 connected on USB.
    If this doesn't work, you'll have to use manually the proxmark3 client, see "$CLIENT -h".
    To see a list of available ports, use --list.

Usage:
    $SCRIPT [-n <N>]
    $SCRIPT --list
EOF
  }
else
  echo >&2 "[!!] Script ran under unknown name, abort: $SCRIPT"
  exit 1
fi

# priority to the help options
for ARG; do
    if [ "$ARG" == "-h" ] || [ "$ARG" == "--help" ]; then
        HELP
        exit 0
    fi
    if [ "$ARG" == "-hh" ] || [ "$ARG" == "--helpclient" ]; then
        CMD "-h"
        exit 0
    fi
done

# if offline, bypass the script and forward all other args
for ARG; do
    shift
    if [ "$ARG" == "-o" ] || [ "$ARG" == "--offline" ]; then
        CMD "$@"
        exit $?
    fi
    set -- "$@" "$ARG"
done

# if a port is already provided, let's just run the command as such
for ARG; do
    if [ "$ARG" == "-p" ]; then
        CMD "$@"
        exit $?
    fi
done

if [ "$1" == "--list" ]; then
    shift
    if [ "$1" != "" ]; then
        echo >&2 "[!!] Option --list must be used alone"
        exit 1
    fi
    SHOWLIST=true
fi

# Number of the proxmark3 we're interested in
N=1
if [ "$1" == "-n" ]; then
    shift
    if [ "$1" -ge 1 ] && [ "$1" -lt 10 ]; then
        N=$1
        shift
    else
        echo >&2 "[!!] Option -n requires a number between 1 and 9, got \"$1\""
        exit 1
    fi
fi

HOSTOS=$(uname | awk '{print toupper($0)}')
if [ "$HOSTOS" = "LINUX" ]; then
    if uname -a|grep -q Microsoft; then
        # First try finding it using the PATH environment variable
        PSHEXE=$(command -v powershell.exe 2>/dev/null)

        # If it fails (such as if WSLENV is not set), try using the default installation path
        if [ -z "$PSHEXE" ]; then
            PSHEXE=/mnt/c/Windows/System32/WindowsPowerShell/v1.0/powershell.exe
        fi

        # Finally test if PowerShell is working
        if ! "$PSHEXE" exit >/dev/null 2>&1; then
            echo >&2 "[!!] Cannot run powershell.exe, are you sure your WSL is authorized to run Windows processes? (cf WSL interop flag)"
            exit 1
        fi

        GETPM3LIST=get_pm3_list_WSL
    else
        GETPM3LIST=get_pm3_list_Linux
    fi
elif [ "$HOSTOS" = "DARWIN" ]; then
    GETPM3LIST=get_pm3_list_macOS
elif [[ "$HOSTOS" =~ MINGW(32|64)_NT* ]]; then
    GETPM3LIST=get_pm3_list_Windows
else
    echo >&2 "[!!] Host OS not recognized, abort: $HOSTOS"
    exit 1
fi

if $SHOWLIST; then
    # Probe for up to 9 devs
    $GETPM3LIST 9
    if [ ${#PM3LIST} -lt 1 ]; then
        echo >&2 "[!!] No port found"
        exit 1
    fi
    n=1
    for DEV in "${PM3LIST[@]}"
        do
            echo "$n: $DEV"
            n=$((n+1))
        done
    exit 0
fi

# Wait till we get at least N proxmark3 devices
$GETPM3LIST "$N"
if [ ${#PM3LIST} -lt "$N" ]; then
    echo >&2 "[=] Waiting for Proxmark3 to appear..."
fi
while true; do
    if [ ${#PM3LIST[*]} -ge "$N" ]; then
        break
    fi
    sleep .1
    $GETPM3LIST "$N"
done

if [ ${#PM3LIST} -lt "$N" ]; then
    HELP() {
      cat << EOF
[!!] No port found, abort

[?] Hint: try '$SCRIPT --list' to see list of available ports,  and use the -n command like below
[?]    $SCRIPT [-n <N>]

EOF
  }
    HELP
    exit 1
fi

CMD "${PM3LIST[$((N-1))]}" "$@"
exit $?