#!/bin/bash

set -e
PROG=$0

function die {
        echo "$PROG: $@" >&2
        exit 1
}

IN=
OUT=
ENTRY=
TMP=$(tempfile)
TMP2=$(tempfile)
TMP3=$(tempfile)

while [ "x${1:0:1}" = "x-" ] ; do
        case $1 in
        -h)
                echo "Disassebler script"
                echo "$0 [-o outfile] [-e entry] binary"
                exit 0
                ;;
        -o)
                shift
                OUT=$1
                ;;
        -e)
                shift
                ENTRY="$ENTRY $1"
                ;;
        -*)
                die "invalid option $1"
                ;;
        esac
        shift
done

IN=$1
[ -n "$IN" ] || die "specify input file"
[ -f "$IN" ] || die "invalid input file: $IN"

ENTRY=${ENTRY:-0}

#echo IN=$IN       >&2
#echo OUT=$OUT     >&2
#echo ENTRY=$ENTRY >&2

# convert hex number to %08x
function fmt_08X {
        printf "%08X" 0x${1/0x/}
}

# this function uses nasm to disassemble up to the next ret or jmp instruction
function do_disasm {
        START=$1
        STOP=$2

        END_LIMIT=
        if [ -n "$STOP" ] ; then
                END_LIMIT="-k $STOP,0xFFFF"
        fi

        FILTER="-e ret -e jmp"

        ( ndisasm -a -b 16 -k 0,$START $END_LIMIT $IN ; echo END) \
        | grep -m1 -B100000 -e ^END $FILTER \
        | grep -v -e ^END -e '00000000  skipping '
}

# this function injects markers with invalid opcode 0000 (that's used for sorting)
# and then calls do_disasm to chin any called or jumpt to code
SPACES="                                                       # "
BADOPCODE="00000000XXXX"
BADOPCODE_MATCH="^[0-9A-Z]\{8\}  $BADOPCODE [^:]*:"
function do_extract {

        grep -v "$BADOPCODE_MATCH" \
        | grep -e "jmp 0xf000:" \
             -e "jmp short 0x" \
             -e "jn\?[a-z]e\? 0x" \
             -e "jz 0x" \
             -e "jnz 0x" \
             -e "loop 0x" \
             -e "call 0x" \
        | while read line ; do 
                #echo x "$line" >&2
                ofs=${line/ */}
                ins=${line:28}
                case "$ins" in
                jmp\ 0xf000:*)
                        HEX=${line/*:/}
                        ADR=$(fmt_08X $HEX)
                        echo "$ADR  $BADOPCODE $ofs: $SPACES long jump 0xf000:$HEX from $ofs"
                        do_disasm $HEX
                        ;;
                jmp\ short\ 0x*)
                        HEX=${line/* /}
                        ADR=$(fmt_08X $HEX)
                        echo "$ADR  $BADOPCODE $ofs: $SPACES short jump $HEX from $ofs"
                        do_disasm $HEX
                        ;;
                loop\ 0x*)
                        HEX=${line/* /}
                        ADR=$(fmt_08X $HEX)
                        echo "$ADR  $BADOPCODE $ofs: $SPACES loop $HEX from $ofs"
                        do_disasm $HEX
                        ;;
                call\ 0x*)
                        HEX=${line/* /}
                        ADR=$(fmt_08X $HEX)
                        echo "$ADR  $BADOPCODE $ofs: $SPACES function $HEX from $ofs"
                        do_disasm $HEX
                        ;;
                j*\ 0x*)
                        HEX=${line/* /}
                        ADR=$(fmt_08X $HEX)
                        echo "$ADR  $BADOPCODE $ofs: $SPACES conditional jump $HEX from $ofs"
                        do_disasm $HEX
                        ;;
                *)
                        echo "UNHANDLED: $line" >&2
                        ;;
                esac
        done
}

# this function removes duplicates and sorts the source back into address order
# if there is no change between the previous iteration it returns 1, otherwise 0
function do_step {

        #echo step >&2

        do_extract < $TMP > $TMP2

        sort -u -t ' ' -k 1,1 -k 2,2 -k 3,3 -k 4,4 <(cat $TMP $TMP2) > $TMP3
        if ( diff "$TMP" "$TMP3" > /dev/null ) ; then
                return 1
        fi
        mv $TMP3 $TMP

        return 0
}

# start with the entry offset(s) given
for ent in $ENTRY ; do
        ADR=$(fmt_08X ${ent/0x/})
        echo "$ADR  $BADOPCODE 0: $SPACES ENTRY $ent" | tr 'a-z' 'A-Z'
        do_disasm 0x$ADR
done > $TMP

# loop as long as we are pulling in new code
while do_step ; do :; done

# remove markers
sed "s/$BADOPCODE_MATCH//" $TMP 

# cleanup
rm -f $TMP $TMP2 $TMP3