; PE-basic (CUI) v0.02
; Developed by Mat-Quasar, 2025/3
;
; Version history: v0.01 - Initial release
;                  v0.02 - Bug fix: Support long filename enclosed in quotes
;
format PE console
include 'win32a.inc'

PE32 = 0x10b
PE64 = 0x20b
IMAGE_SUBSYSTEM_WINDOWS_GUI = 2
IMAGE_SUBSYSTEM_WINDOWS_CUI = 3
Offset_MagicNumber = 0x98 - 0x80
Offset_Subsystem = 0xDC - 0x80

section '.code' code readable executable

  entry $
        invoke  GetStdHandle, STD_OUTPUT_HANDLE
        mov     dword [_stdout],eax

        call    [GetCommandLine]
        push    eax
        cld
        mov     edi, eax
        or      ecx, -1
        xor     eax, eax
        repnz   scasb           ; Calculate total length of command line arguments
        not     ecx
        pop     eax
        push    eax
        mov     dword [_fnlen], ecx
        mov     edi, eax
        or      ecx, -1
        mov     eax, 32
        repnz   scasb           ; Find the first white space after first argument
        repz    scasb           ; Scan for extra white spacing between first argument and second argument
        not     ecx
        mov     dword [_dummy], ecx
        sub     dword [_fnlen], ecx
        cmp     dword [_fnlen], 0    ; Check if no argument is supplied
        jle     .err0
        mov     ecx, dword [_fnlen]
        pop     eax
        add     eax, dword [_dummy]
        dec     eax
        xor     ebx, ebx
        cmp     byte [eax], '"'      ; Check if path is enclosed in quotes
        setz    bl
        sub     ecx, ebx
        sub     ecx, ebx
        mov     esi, eax
        add     esi, ebx
        lea     edi, [_filename]
        rep     movsb

        invoke  WriteFile, dword [_stdout], _filename, dword [_fnlen], _dummy, 0

        invoke  CreateFile, _filename, GENERIC_READ, FILE_SHARE_READ, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0
        mov     dword [_handle], eax
        cmp     eax, INVALID_HANDLE_VALUE
        je      .err1

        invoke  ReadFile, dword [_handle], _MZ_ID_r, 2, _dummy, 0
        movzx   ebx, word [_MZ_ID_r]
        cmp     bx, word [_MZ_ID]
        jnz     .err4
        invoke  SetFilePointer, dword [_handle], 0x3C, 0, FILE_BEGIN
        invoke  ReadFile, dword [_handle], _PE_start, 4, _dummy, 0
        invoke  SetFilePointer, dword [_handle], dword [_PE_start], 0, FILE_BEGIN
        invoke  ReadFile, dword [_handle], _PE_ID_r, 4, _dummy, 0
        mov     ebx, dword [_PE_ID_r]
        cmp     ebx, dword [_PE_ID]
        jnz     .err5
        mov     ebx, dword [_PE_start]
        add     ebx, Offset_MagicNumber
        mov     [_offset], ebx
        invoke  SetFilePointer, dword [_handle], dword [_offset] , 0, FILE_BEGIN
        invoke  ReadFile, dword [_handle], _magic, 2, _dummy, 0
        movzx   ebx, word [_magic]
        cmp     ebx, PE32
        jz      .isPE32
        invoke  WriteFile, dword [_stdout], _msg7, _msglen7, _dummy, 0
        jmp     .continue
  .isPE32:
        invoke  WriteFile, dword [_stdout], _msg6, _msglen6, _dummy, 0
  .continue:

        mov     ebx, dword [_PE_start]
        add     ebx, Offset_Subsystem
        mov     [_offset], ebx
        invoke  SetFilePointer, dword [_handle], dword [_offset] , 0, FILE_BEGIN
        invoke  ReadFile, dword [_handle], _subsystem, 2, _dummy, 0
        movzx   ebx, word [_subsystem]
        cmp     ebx, IMAGE_SUBSYSTEM_WINDOWS_CUI
        jz      .isCUI
        cmp     ebx, IMAGE_SUBSYSTEM_WINDOWS_GUI
        jz      .isGUI
        invoke  WriteFile, dword [_stdout], _msg10, _msglen10, _dummy, 0
        jmp     .done
  .isCUI:
        invoke  WriteFile, dword [_stdout], _msg8, _msglen8, _dummy, 0
        jmp     .done
  .isGUI:
        invoke  WriteFile, dword [_stdout], _msg9, _msglen9, _dummy, 0
  .done:
        invoke  CloseHandle, dword [_handle]
        jmp     .quit

  .err0:
        invoke  WriteFile, dword [_stdout], _msg0, _msglen0, _dummy, 0
        jmp     .quit
  .err1:
        invoke  WriteFile, dword [_stdout], _msg1, _msglen1, _dummy, 0
        jmp     .quit
  .err2:
        invoke  WriteFile, dword [_stdout], _msg2, _msglen2, _dummy, 0
        jmp     .quit
  .err4:
        invoke  CloseHandle, dword [_handle]
        invoke  WriteFile, dword [_stdout], _msg4, _msglen4, _dummy, 0
        jmp     .quit
  .err5:
        invoke  CloseHandle, dword [_handle]
        invoke  WriteFile, dword [_stdout], _msg5, _msglen5, _dummy, 0
        jmp     .quit

  .quit:
        invoke  WriteFile, dword [_stdout], _msg3, _msglen3, _dummy, 0
        invoke  GetStdHandle, STD_INPUT_HANDLE
        invoke  ReadConsole, eax, _key, 1, _dummy, 0
        invoke  ExitProcess, 0

section '.data' data readable writable

  _msg0 db 'Usage: pebasic <filename>',13,10
  _msglen0  = $ - _msg0
  _msg1 db 13,10,'Error opening file.',13,10
  _msglen1  = $ - _msg1
  _msg2 db 13,10,'Error reading file.',13,10
  _msglen2  = $ - _msg2
  _msg3 db 13,10,'Press Enter to quit...',13,10
  _msglen3  = $ - _msg3
  _msg4 db 13,10,'Invalid EXE.',13,10
  _msglen4  = $ - _msg4
  _msg5 db 13,10,'Invalid PE.',13,10
  _msglen5  = $ - _msg5
  _msg6 db 13,10,'32-bit '
  _msglen6 = $ - _msg6
  _msg7 db 13,10,'64-bit '
  _msglen7 = $ - _msg7
  _msg8 db 'Console App.',13,10
  _msglen8 = $ - _msg8
  _msg9 db 'GUI App.',13,10
  _msglen9 = $ - _msg9
  _msg10 db 'PE file.',13,10
  _msglen10 = $ - _msg10


  _fnlen    dd ?
  _subsystem dd ?
  _dummy    dd ?
  _handle   dd ?
  _stdout   dd ?
  _key      dd ?
  _offset   dd 0
  _PE_start dd ?
  _PE_ID_r  rb 4
  _PE_ID    db 'P','E',0,0
  _MZ_ID_r  rb 2
  _MZ_ID    db 'M','Z'
  _magic    dw ?
  _filename rb MAX_PATH

section '.idata' import readable writable

 library kernel32, 'KERNEL32.DLL'

 import kernel32,\
        GetStdHandle, 'GetStdHandle', \
        WriteConsole, 'WriteConsoleA', \
        ReadConsole, 'ReadConsoleA', \
        CreateFile, 'CreateFileA', \
        ReadFile, 'ReadFile', \
        WriteFile, 'WriteFile', \
        CloseHandle, 'CloseHandle', \
        GetCommandLine, 'GetCommandLineA', \
        SetFilePointer, 'SetFilePointer', \
        ExitProcess, 'ExitProcess'