/*! @file       i386/boot1pxe/boot.c
    @copyright  2007,2009 David F. Elliott
    @abstract   Contains main booter function of "remote.0" style PXE booter
 */
#include "libsaio.h"

#include <stdint.h>

#include "pxe.h"

typedef struct b1pxe0_asm
{   uint8_t     jmp_instructions[3];
    uint8_t     mov_ax_opcode;
    uint16_t    mov_ax_operand;
} __attribute__((packed)) b1pxe0_asm;

static void zeroBSS()
{
    extern struct {}    _DATA__bss__begin ASM_MACHO_SECTION_BEGIN(__DATA,__bss),
                        _DATA__bss__end ASM_MACHO_SECTION_END(__DATA,__bss);
    extern struct {}    _DATA__common__begin ASM_MACHO_SECTION_BEGIN(__DATA,__common),
                        _DATA__common__end ASM_MACHO_SECTION_END(__DATA,__common);

    // Save the original EBP before we zero BSS
    uint32_t real_entry_ebp = gPXE_real_entry_ebp;

    bzero( &_DATA__bss__begin,
           ((char *)&_DATA__bss__end - (char *)&_DATA__bss__begin) );

    bzero( &_DATA__common__begin, 
           ((char *)&_DATA__common__end - (char *)&_DATA__common__begin) );

    // Restore the original EBP now that we've zeroed BSS
    gPXE_real_entry_ebp = real_entry_ebp;
}

#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wunused-function"
static void sleepCountdown(int secs) {
    int i;
    for(i=secs; i>0; --i)
    {
        printf("%d...", i);
        sleep(1);
    }
    printf("0\n");
}
#pragma clang diagnostic pop

extern SEGOFF16 pxenv_segoff;
extern SEGOFF16 pxe_segoff;

char const BOOT2_FILENAME[] = "boot";

uint16_t boot(uint32_t biosdev)
{
    zeroBSS();
    // NOTE: Don't enable A20

    printf("Darwin/x86 first-stage NBP\n");
    
    printf("'PXENV+' should be at %04x:%04xh\n", pxenv_segoff.segment, pxenv_segoff.offset);
    printf("'!PXE' should be at %04x:%04xh\n", pxe_segoff.segment, pxe_segoff.offset);
    
    initPxeFromBootstrap(
        (void*)ADDR32(pxe_segoff.segment, pxe_segoff.offset),
        (void*)ADDR32(pxenv_segoff.segment, pxenv_segoff.offset));

#if 0
    PXEStartup();
#endif

    //sleep5s();

    uint16_t pxenv_exit;
    t_PXENV_GET_CACHED_INFO cachedInfoData;
    BOOTPLAYER *dhcp_discover_packet = NULL;
    BOOTPLAYER const *dhcp_ack_packet = NULL;
    size_t dhcp_ack_packet_sz = 0;
    printf("Retrieving DHCP Discover info\n");
    if((pxenv_exit = pxenvGetCachedInfo(&cachedInfoData, PXENV_PACKET_TYPE_DHCP_DISCOVER)) == PXENV_EXIT_SUCCESS)
        dhcp_discover_packet = getPxeCachedBuffer(&cachedInfoData);
#if 1
    printf("opcode\t%d\n", dhcp_discover_packet->opcode);
    printf("MAC %02x:%02x:%02x:%02x:%02x:%02x\n"
        ,   dhcp_discover_packet->CAddr[0]
        ,   dhcp_discover_packet->CAddr[1]
        ,   dhcp_discover_packet->CAddr[2]
        ,   dhcp_discover_packet->CAddr[3]
        ,   dhcp_discover_packet->CAddr[4]
        ,   dhcp_discover_packet->CAddr[5]
      );
#endif

    printf("Retrieving DHCP ACK info\n");
#if 1
    if((pxenv_exit = pxenvGetCachedInfo(&cachedInfoData, PXENV_PACKET_TYPE_DHCP_ACK)) == PXENV_EXIT_SUCCESS) {
        dhcp_ack_packet = getPxeCachedBuffer(&cachedInfoData);
        dhcp_ack_packet_sz = cachedInfoData.BufferSize;
    }
    printf("Client IP %d.%d.%d.%d\n"
        ,   dhcp_ack_packet->cip.array[0]
        ,   dhcp_ack_packet->cip.array[1]
        ,   dhcp_ack_packet->cip.array[2]
        ,   dhcp_ack_packet->cip.array[3]
      );

    printf("'Your' IP %d.%d.%d.%d\n"
        ,   dhcp_ack_packet->yip.array[0]
        ,   dhcp_ack_packet->yip.array[1]
        ,   dhcp_ack_packet->yip.array[2]
        ,   dhcp_ack_packet->yip.array[3]
      );

    printf("Server IP %d.%d.%d.%d\n"
        ,   dhcp_ack_packet->sip.array[0]
        ,   dhcp_ack_packet->sip.array[1]
        ,   dhcp_ack_packet->sip.array[2]
        ,   dhcp_ack_packet->sip.array[3]
      );

    printf("Gateway IP %d.%d.%d.%d\n"
        ,   dhcp_ack_packet->gip.array[0]
        ,   dhcp_ack_packet->gip.array[1]
        ,   dhcp_ack_packet->gip.array[2]
        ,   dhcp_ack_packet->gip.array[3]
      );
#endif

    // NOTE: Depending on BIOS executing this next call may invalidate the previous buffer.
    // I have nothing to back this up I just would suspect as much out of a BIOS call.
    printf("Retrieving Boot Server DHCP reply\n");
    if((pxenv_exit = pxenvGetCachedInfo(&cachedInfoData, PXENV_PACKET_TYPE_CACHED_REPLY)) == PXENV_EXIT_SUCCESS) {
        dhcp_ack_packet = getPxeCachedBuffer(&cachedInfoData);
        dhcp_ack_packet_sz = cachedInfoData.BufferSize;
    }

#if 1
    printf("Server IP %d.%d.%d.%d\n"
        ,   dhcp_ack_packet->sip.array[0]
        ,   dhcp_ack_packet->sip.array[1]
        ,   dhcp_ack_packet->sip.array[2]
        ,   dhcp_ack_packet->sip.array[3]
      );
#endif
    // The typical number of options is about a half dozen so keeping the array of
    // pointers to them on the stack is reasonable.  The absolute highest number
    // of recorded options is frame_size (typically 1500) - 20 (IP header) - 8 (UDP)
    // - 236 (size of BOOTP up to options) - 4 (size of DHCP magic) which is 1232
    // divided by 2 which is 616.  For example, someone could use option 43 with
    // length 0 616 times.  While it's theoretically possible to smash the stack
    // into something by doing this it'd be a damn difficult attack.
    size_t option_list_sz = parseOptionsFromDhcpPacket(NULL, 0, dhcp_ack_packet, dhcp_ack_packet_sz);
    unsigned char const *option_list[option_list_sz];

    // Read in all of the options
    size_t option_list_len = 0;
    option_list_len = parseOptionsFromDhcpPacket(option_list, option_list_sz, dhcp_ack_packet, dhcp_ack_packet_sz);
    if (option_list_len > option_list_sz) {
        printf("Option list length greater?");
        return 1;
    }

    // Print out all options for debugging.
    debugPrintOptions(option_list, 0, option_list_len, dhcp_ack_packet);
    // sleepCountdown(3);

    t_PXENV_TFTP_GET_FSIZE tftpGetFsizeCommand;
    bzero(&tftpGetFsizeCommand, sizeof(tftpGetFsizeCommand));

    // NOTE: Although there is sname and possibly option 66 the PXE API doesn't
    // support DNS lookups nor do them internally.  The only valid way to contact
    // the TFTP server is by using siaddr containing the IPv4 address of the
    // next boot server which must handle TFTP.
    // The confusion here arises because MS DHCP actually allows you to specify
    // option 66 which it resolves on the server side causing siaddr to be
    // set to the resolved IP.
    // The equivalent with ISC DHCP is to set both option 66 and siaddr
    // which are the tftp-server-name and next-server options.
    // Note that this is compatible with the RFC 2131 commentary that
    // DHCP clarifies that siaddr is indeed the IP address of the next server
    // and not of the DHCP server itself (that is option 54)
    tftpGetFsizeCommand.ServerIPAddress = dhcp_ack_packet->sip;

    // Read the base file (the portion before the last slash, if any)
    // into the command buffer
    size_t baseFilenameLen = setBaseFilenameFromOptions(
            tftpGetFsizeCommand.FileName, sizeof(tftpGetFsizeCommand.FileName),
            dhcp_ack_packet, option_list, option_list_len);

    // Concatenate that with "boot", the name of the boot2 file
    char *out = tftpGetFsizeCommand.FileName + baseFilenameLen;
    char const *in;

    if(out + sizeof(BOOT2_FILENAME) < tftpGetFsizeCommand.FileName + sizeof(tftpGetFsizeCommand.FileName))
    {
        for(in = BOOT2_FILENAME; in < BOOT2_FILENAME + sizeof(BOOT2_FILENAME); ++in, ++out)
            *out = *in;
    }
    else
    {
        printf("Computed boot filename is too long.  Consider shortening the path.\n");
        return -1;
    }

    printf("Getting size of TFTP file \"%s\"\n", tftpGetFsizeCommand.FileName);

    pxenv_exit = pxeCall(PXENV_TFTP_GET_FSIZE, &tftpGetFsizeCommand);
    if(pxenv_exit != PXENV_EXIT_SUCCESS)
    {
        printf("TFTP Get Fsize call failed with exit %d\n", pxenv_exit);
        sleep(4);
        return -1;
    }

    if(tftpGetFsizeCommand.Status != PXENV_STATUS_SUCCESS)
    {
        printf("TFTP Get Fsize failed with status %d\n", tftpGetFsizeCommand.Status);
        sleep(4);
        return -1;
    }
    printf("TFTP returned file size of %ld bytes\n", tftpGetFsizeCommand.FileSize);
    //sleep5s();

    

    //////////////////////////////////////////////////////

    printf("Opening TFTP file \"%s\"\n", tftpGetFsizeCommand.FileName);
    t_PXENV_TFTP_OPEN tftpOpenCommand;
    bzero(&tftpOpenCommand, sizeof(tftpOpenCommand));

    // FIXME: TFTP server can be different and comes from DHCP options
    tftpOpenCommand.ServerIPAddress = dhcp_ack_packet->sip;
    strcpy(tftpOpenCommand.FileName, tftpGetFsizeCommand.FileName);
    tftpOpenCommand.TFTPPort = 69 << 8; // byte-swapped.. yeah.. it's in network byte order!
    tftpOpenCommand.PacketSize = 1440;
    
    pxenv_exit = pxeCall(PXENV_TFTP_OPEN, &tftpOpenCommand);
    if(pxenv_exit != PXENV_EXIT_SUCCESS)
    {
        printf("TFTP Open call failed with exit %d\n", pxenv_exit);
        sleep(4);
        return -1;
    }

    if(tftpOpenCommand.Status != PXENV_STATUS_SUCCESS)
    {
        printf("TFTP Open failed with status %d\n", tftpOpenCommand.Status);
        sleep(4);
        return -1;
    }

    printf("Suggested packet size is %d bytes\n", tftpOpenCommand.PacketSize);
    //sleep5s();

    //////////////////////////////////////////////////////

    printf("We would transfer the file now.\n");

    t_PXENV_TFTP_READ tftpReadCommand;
    bzero(&tftpReadCommand, sizeof(tftpReadCommand));

    uint32_t BytesRemaining = tftpGetFsizeCommand.FileSize;
    char *Buffer = (char*)(uint32_t)(BOOT2_ADDR);
        // Make buffer size 32k which should be a safe value
        tftpReadCommand.BufferSize = 32 * 1024;
    for(;BytesRemaining > 0;)
    {
        tftpReadCommand.Buffer.segment = NORMALIZED_SEGMENT((uint32_t)Buffer);
        tftpReadCommand.Buffer.offset = NORMALIZED_OFFSET((uint32_t)Buffer);

        // BufferSize is set to the amount read each time through.

        pxenv_exit = pxeCall(PXENV_TFTP_READ, &tftpReadCommand);
        if(pxenv_exit != PXENV_EXIT_SUCCESS)
        {
            printf("TFTP Read call failed with exit %d\n", pxenv_exit);
            sleep(4);
            return -1;
        }
        if(tftpReadCommand.Status != PXENV_STATUS_SUCCESS)
        {
            printf("TFTP Read failed with status %d\n", tftpOpenCommand.Status);
            sleep(4);
            return -1;
        }
        BytesRemaining -= tftpReadCommand.BufferSize;
        Buffer += tftpReadCommand.BufferSize;
    }

    //////////////////////////////////////////////////////

    tftpOpenCommand.Status = 0;
    pxenv_exit = pxeCall(PXENV_TFTP_CLOSE, &tftpOpenCommand);
    if(pxenv_exit != PXENV_EXIT_SUCCESS)
    {
        printf("TFTP Close failed with exit %d\n", pxenv_exit);
        return -1;
    }
    if(tftpOpenCommand.Status != PXENV_STATUS_SUCCESS)
    {
        printf("TFTP Close failed with status %d\n", tftpOpenCommand.Status);
        return -1;
    }

    printf("Seemed to close okay!\n");
    // sleepCountdown(3);

    return 0;
}