/*

  *  Copyright (C) 2014-2019 Cisco Systems, Inc. and/or its affiliates. All rights reserved.

  *
  *  Authors: Kevin Lin <klin@sourcefire.com>
  *
  *  This program is free software; you can redistribute it and/or modify
  *  it under the terms of the GNU General Public License version 2 as
  *  published by the Free Software Foundation.
  *
  *  This program is distributed in the hope that it will be useful,
  *  but WITHOUT ANY WARRANTY; without even the implied warranty of
  *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  *  GNU General Public License for more details.
  *
  *  You should have received a copy of the GNU General Public License
  *  along with this program; if not, write to the Free Software
  *  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
  *  MA 02110-1301, USA.
  */
 
 #if HAVE_CONFIG_H
 #include "clamav-config.h"
 #endif
 
 #include <stdio.h>
 #include <errno.h>
 #if HAVE_STRING_H
 #include <string.h>
 #endif
 #include <ctype.h>
 #include <fcntl.h>
 #include <zlib.h>
 

 #include "clamav.h"

 #include "others.h"
 #include "mbr.h"
 #include "prtn_intxn.h"
 #include "scanners.h"

 #include "dconf.h"

 
 //#define DEBUG_MBR_PARSE
 //#define DEBUG_EBR_PARSE
 

 #ifndef PRTN_INTXN_DETECTION
 #  define PRTN_INTXN_DETECTION "heuristic.mbrprtnintersect"
 #endif
 

 #ifdef DEBUG_MBR_PARSE
 #  define mbr_parsemsg(...) cli_dbgmsg( __VA_ARGS__)
 #else
 #  define mbr_parsemsg(...) ;
 #endif
 
 #ifdef DEBUG_EBR_PARSE
 #  define ebr_parsemsg(...) cli_dbgmsg( __VA_ARGS__)
 #else
 #  define ebr_parsemsg(...) ;
 #endif
 
 enum MBR_STATE {
     SEEN_NOTHING,
     SEEN_PARTITION,
     SEEN_EXTENDED,
     SEEN_EMPTY
 };
 

 static int mbr_scanextprtn(cli_ctx *ctx, unsigned *prtncount, off_t extlba, 
                            size_t extlbasize, size_t sectorsize);
 static int mbr_check_mbr(struct mbr_boot_record *record, size_t maplen, size_t sectorsize);

 static int mbr_check_ebr(struct mbr_boot_record *record);
 static int mbr_primary_prtn_intxn(cli_ctx *ctx, struct mbr_boot_record mbr, size_t sectorsize);

 static int mbr_extended_prtn_intxn(cli_ctx *ctx, unsigned *prtncount, off_t extlba, size_t sectorsize);

 
 int cli_mbr_check(const unsigned char *buff, size_t len, size_t maplen) {
     struct mbr_boot_record mbr;

     off_t mbr_base = 0;
     size_t sectorsize = 512;

     if (len < sectorsize) {

         return CL_EFORMAT;
     }
 

     mbr_base = sectorsize - sizeof(struct mbr_boot_record);
     memcpy(&mbr, buff+mbr_base, sizeof(mbr));

     mbr_convert_to_host(&mbr);
 

     if ((mbr.entries[0].type == MBR_PROTECTIVE) || (mbr.entries[0].type == MBR_HYBRID))
         return CL_TYPE_GPT;

     return mbr_check_mbr(&mbr, maplen, sectorsize);

 }
 

 int cli_mbr_check2(cli_ctx *ctx, size_t sectorsize) {

     struct mbr_boot_record mbr;

     off_t pos = 0, mbr_base = 0;
     size_t maplen;

     if (!ctx || !ctx->fmap) {
         cli_errmsg("cli_scanmbr: Invalid context\n");
         return CL_ENULLARG;

     }
 

     /* sector size calculation, actual value is OS dependent */
     if (sectorsize == 0)
         sectorsize = MBR_SECTOR_SIZE;
 

     mbr_base = sectorsize - sizeof(struct mbr_boot_record);
 

     /* size of total file must be a multiple of the sector size */
     maplen = (*ctx->fmap)->real_len;
     if ((maplen % sectorsize) != 0) {

         cli_dbgmsg("cli_scanmbr: File sized %lu is not a multiple of sector size %lu\n",
                    (unsigned long)maplen, (unsigned long)sectorsize);

         return CL_EFORMAT;
     }
 
     /* sector 0 (first sector) is the master boot record */
     pos = (MBR_SECTOR * sectorsize) + mbr_base;
 
     /* read the master boot record */
     if (fmap_readn(*ctx->fmap, &mbr, pos, sizeof(mbr)) != sizeof(mbr)) {
         cli_dbgmsg("cli_scanmbr: Invalid master boot record\n");
         return CL_EFORMAT;
     }
 

     /* convert the little endian to host, include the internal  */
     mbr_convert_to_host(&mbr);
 

     if ((mbr.entries[0].type == MBR_PROTECTIVE) || (mbr.entries[0].type == MBR_HYBRID))
         return CL_TYPE_GPT;
 
     return mbr_check_mbr(&mbr, maplen, sectorsize);

 }
 

 /* sets sectorsize to default value if specified to be 0 */

 int cli_scanmbr(cli_ctx *ctx, size_t sectorsize)

 {
     struct mbr_boot_record mbr;
     enum MBR_STATE state = SEEN_NOTHING;

     int ret = CL_CLEAN, detection = CL_CLEAN;

     off_t pos = 0, mbr_base = 0, partoff = 0;

     unsigned i = 0, prtncount = 0;

     size_t maplen, partsize;

 
     mbr_parsemsg("The start of something magnificant: MBR parsing\n");
 
     if (!ctx || !ctx->fmap) {
         cli_errmsg("cli_scanmbr: Invalid context\n");
         return CL_ENULLARG;
     }
 

     /* sector size calculation, actual value is OS dependent */

     if (sectorsize == 0)
         sectorsize = MBR_SECTOR_SIZE;
 
     mbr_base = sectorsize - sizeof(struct mbr_boot_record);

 
     /* size of total file must be a multiple of the sector size */
     maplen = (*ctx->fmap)->real_len;
     if ((maplen % sectorsize) != 0) {

         cli_dbgmsg("cli_scanmbr: File sized %lu is not a multiple of sector size %lu\n",
                    (unsigned long)maplen, (unsigned long)sectorsize);

         return CL_EFORMAT;
     }
 
     /* sector 0 (first sector) is the master boot record */

     pos = (MBR_SECTOR * sectorsize) + mbr_base;

 
     /* read the master boot record */
     if (fmap_readn(*ctx->fmap, &mbr, pos, sizeof(mbr)) != sizeof(mbr)) {
         cli_dbgmsg("cli_scanmbr: Invalid master boot record\n");
         return CL_EFORMAT;
     }
 
     /* convert the little endian to host, include the internal  */
     mbr_convert_to_host(&mbr);
 
     /* MBR checks */

     ret = mbr_check_mbr(&mbr, maplen, sectorsize);

     if (ret != CL_CLEAN) {
         return ret;
     }
 

     /* MBR is valid, examine bootstrap code */

     ret = cli_map_scan(*ctx->fmap, 0, sectorsize, ctx, CL_TYPE_ANY);

     if (ret != CL_CLEAN) {

         if (SCAN_ALLMATCHES && (ret == CL_VIRUS))

             detection = CL_VIRUS;
         else
             return ret;

     }
 

     /* check that the partition table has no intersections - HEURISTICS */

     if (SCAN_HEURISTIC_PARTITION_INTXN && (ctx->dconf->other & OTHER_CONF_PRTNINTXN)) {

         ret = mbr_primary_prtn_intxn(ctx, mbr, sectorsize);

         if (ret != CL_CLEAN) {

             if (SCAN_ALLMATCHES && (ret == CL_VIRUS))

                 detection = CL_VIRUS;
             else
                 return ret;

         }
     }
 
     /* MBR is valid, examine partitions */
     prtncount = 0;
     cli_dbgmsg("MBR Signature: %x\n", mbr.signature);

     for (i = 0; i < MBR_MAX_PARTITION_ENTRIES && prtncount < ctx->engine->maxpartitions; ++i) {

         cli_dbgmsg("MBR Partition Entry %u:\n", i);
         cli_dbgmsg("Status: %u\n", mbr.entries[i].status);
         cli_dbgmsg("Type: %x\n", mbr.entries[i].type);

         cli_dbgmsg("Blocks: [%u, +%u), ([%lu, +%lu))\n",

                    mbr.entries[i].firstLBA, mbr.entries[i].numLBA,

                    (unsigned long)(mbr.entries[i].firstLBA * sectorsize),
                    (unsigned long)(mbr.entries[i].numLBA * sectorsize));

 
         /* Handle MBR entry based on type */
         if (mbr.entries[i].type == MBR_EMPTY) {

             /* empty partition entry */

             prtncount++;
         }
         else if (mbr.entries[i].type == MBR_EXTENDED) {
             if (state == SEEN_EXTENDED) {
                 cli_dbgmsg("cli_scanmbr: detected a master boot record "
                            "with multiple extended partitions\n");
             }

             state = SEEN_EXTENDED; /* used only to detect multiple extended partitions */

             ret = mbr_scanextprtn(ctx, &prtncount, mbr.entries[i].firstLBA, 
                                   mbr.entries[i].numLBA, sectorsize);

             if (ret != CL_CLEAN) {

                 if (SCAN_ALLMATCHES && (ret == CL_VIRUS))

                     detection = CL_VIRUS;
                 else
                     return ret;

             }

         }
         else {
             prtncount++;
 
             partoff = mbr.entries[i].firstLBA * sectorsize;

             partsize = mbr.entries[i].numLBA * sectorsize;

             mbr_parsemsg("cli_map_scan: [%u, +%u)\n", partoff, partsize);
             ret = cli_map_scan(*ctx->fmap, partoff, partsize, ctx, CL_TYPE_PART_ANY);

             if (ret != CL_CLEAN) {

                 if (SCAN_ALLMATCHES && (ret == CL_VIRUS))

                     detection = CL_VIRUS;
                 else

                     return ret;
             }
         }
     }
 

     if (prtncount >= ctx->engine->maxpartitions) {
         cli_dbgmsg("cli_scanmbr: maximum partitions reached\n");
     }

 
     return detection;

 }
 

 static int mbr_scanextprtn(cli_ctx *ctx, unsigned *prtncount, off_t extlba, size_t extlbasize, size_t sectorsize)
 {

     struct mbr_boot_record ebr;
     enum MBR_STATE state = SEEN_NOTHING;

     int ret = CL_CLEAN, detection = CL_CLEAN;

     off_t pos = 0, mbr_base = 0, logiclba = 0, extoff = 0, partoff = 0;

     size_t partsize, extsize;
     unsigned i = 0, j = 0;
 
     ebr_parsemsg("The start of something exhausting: EBR parsing\n");
 

     mbr_base = sectorsize - sizeof(struct mbr_boot_record);
 

     logiclba = 0;
     extoff = extlba * sectorsize;
     extsize = extlbasize * sectorsize;
     do {
         pos = extlba * sectorsize; /* start of extended partition */
 
         /* read the extended boot record */

         pos += (logiclba * sectorsize) + mbr_base;

         if (fmap_readn(*ctx->fmap, &ebr, pos, sizeof(ebr)) != sizeof(ebr)) {
             cli_dbgmsg("cli_scanebr: Invalid extended boot record\n");
             return CL_EFORMAT;
         }

         /* convert the little endian to host */
         mbr_convert_to_host(&ebr);

         /* EBR checks */
         ret = mbr_check_ebr(&ebr);
         if (ret != CL_CLEAN) {
             return ret;
         }

         /* update state */
         state = SEEN_NOTHING;
         (*prtncount)++;
 
         /* EBR is valid, examine partitions */
         cli_dbgmsg("EBR Partition Entry %u:\n", i++);
         cli_dbgmsg("EBR Signature: %x\n", ebr.signature);
         for (j = 0; j < MBR_MAX_PARTITION_ENTRIES; ++j) {
             if (j < 2) {
                 cli_dbgmsg("Logical Partition Entry %u:\n", j);
                 cli_dbgmsg("Status: %u\n", ebr.entries[j].status);
                 cli_dbgmsg("Type: %x\n", ebr.entries[j].type);

                 cli_dbgmsg("Blocks: [%u, +%u), ([%lu, +%lu))\n",

                            ebr.entries[j].firstLBA, ebr.entries[j].numLBA,

                            (unsigned long)(ebr.entries[j].firstLBA * sectorsize),
                            (unsigned long)(ebr.entries[j].numLBA * sectorsize));

 
                 if (ebr.entries[j].type == MBR_EMPTY) {

                     /* empty partition entry */

                     switch(state) {
                     case SEEN_NOTHING:
                         state = SEEN_EMPTY;
                         break;
                     case SEEN_PARTITION:
                         logiclba = 0;
                         break;
                     case SEEN_EMPTY:
                         logiclba = 0;
                         /* fall-through */
                     case SEEN_EXTENDED:

                         cli_warnmsg("cli_scanebr: detected a logical boot record "
                                     "without a partition record\n");

                         break;
                     default:

                         cli_warnmsg("cli_scanebr: undefined state for EBR parsing\n");

                         return CL_EPARSE;
                     }

                 }

                 else if (ebr.entries[j].type == MBR_EXTENDED) {
                     switch(state) {
                     case SEEN_NOTHING:
                         state = SEEN_EXTENDED;
                         break;
                     case SEEN_PARTITION:
                         break;
                     case SEEN_EMPTY:

                         cli_warnmsg("cli_scanebr: detected a logical boot record "
                                     "without a partition record\n");

                         break;
                     case SEEN_EXTENDED:

                         cli_warnmsg("cli_scanebr: detected a logical boot record "
                                     "with multiple extended partition records\n");

                         return CL_EFORMAT;
                     default:
                         cli_dbgmsg("cli_scanebr: undefined state for EBR parsing\n");
                         return CL_EPARSE;
                     }

                     logiclba = ebr.entries[j].firstLBA;

                 }

                 else {
                     switch(state) {
                     case SEEN_NOTHING:
                         state = SEEN_PARTITION;
                         break;
                     case SEEN_PARTITION:

                         cli_warnmsg("cli_scanebr: detected a logical boot record "
                                     "with multiple partition records\n");

                         logiclba = 0; /* no extended partitions are possible */
                         break;
                     case SEEN_EXTENDED:

                         cli_warnmsg("cli_scanebr: detected a logical boot record "
                                     "with extended partition record first\n");

                         break;
                     case SEEN_EMPTY:

                         cli_warnmsg("cli_scanebr: detected a logical boot record "
                                     "with empty partition record first\n");

                         logiclba = 0; /* no extended partitions are possible */
                         break;
                     default:
                         cli_dbgmsg("cli_scanebr: undefined state for EBR parsing\n");
                         return CL_EPARSE;
                     }

                     partoff = (extlba + logiclba + ebr.entries[j].firstLBA) * sectorsize;
                     partsize = ebr.entries[j].numLBA * sectorsize;

                     if (partoff + partsize > extoff + extsize) {
                         cli_dbgmsg("cli_scanebr: Invalid extended partition entry\n");

                         return CL_EFORMAT;
                     }
 
                     ret = cli_map_scan(*ctx->fmap, partoff, partsize, ctx, CL_TYPE_PART_ANY);

                     if (ret != CL_CLEAN) {

                         if (SCAN_ALLMATCHES && (ret == CL_VIRUS))

                             detection = CL_VIRUS;
                         else
                             return ret;

                     }

                 }
             }

             else {
                 /* check the last two entries to be empty */
                 if (ebr.entries[j].type != MBR_EMPTY) {
                     cli_dbgmsg("cli_scanebr: detected a non-empty partition "
                                "entry at index %u\n", j);
                     /* should we attempt to use these entries? */
                     return CL_EFORMAT;
                 }

             }
         }

     } while (logiclba != 0 && (*prtncount) < ctx->engine->maxpartitions);

     cli_dbgmsg("cli_scanmbr: examined %u logical partitions\n", i);

 
     return detection;

 }
 

 void mbr_convert_to_host(struct mbr_boot_record *record)
 {
     struct mbr_partition_entry *entry;
     unsigned i;
  
     for (i = 0; i < MBR_MAX_PARTITION_ENTRIES; ++i) {
         entry = &record->entries[i];
  
         entry->firstLBA = le32_to_host(entry->firstLBA);
         entry->numLBA = le32_to_host(entry->numLBA);
     }
     record->signature = be16_to_host(record->signature);
 }
 

 static int mbr_check_mbr(struct mbr_boot_record *record, size_t maplen, size_t sectorsize)

 {
     unsigned i = 0;
     off_t partoff = 0;

     size_t partsize = 0;

 
     for (i = 0; i < MBR_MAX_PARTITION_ENTRIES; ++i) {
         /* check status */
         if ((record->entries[i].status != MBR_STATUS_INACTIVE) && 
             (record->entries[i].status != MBR_STATUS_ACTIVE)) {
             cli_dbgmsg("cli_scanmbr: Invalid boot record status\n");
             return CL_EFORMAT;
         }
 
         partoff = record->entries[i].firstLBA * sectorsize;

         partsize = record->entries[i].numLBA * sectorsize;

         if (partoff + partsize > maplen) {
             cli_dbgmsg("cli_scanmbr: Invalid partition entry\n");
             return CL_EFORMAT;
         }
     }
 
     /* check the signature */
     if (record->signature != MBR_SIGNATURE) {
         cli_dbgmsg("cli_scanmbr: Invalid boot record signature\n");
         return CL_EFORMAT;
     }
 

     /* check the maplen */
     if ((maplen / sectorsize) < 2) {

         cli_dbgmsg("cli_scanmbr: bootstrap code or file is too small to hold disk image\n");

         return CL_EFORMAT;
     }
 

     return CL_CLEAN;
 }
 
 static int mbr_check_ebr(struct mbr_boot_record *record)
 {
     unsigned i = 0;
 

     for (i = 0; i < MBR_MAX_PARTITION_ENTRIES-2; ++i) {

         /* check status */
         if ((record->entries[i].status != MBR_STATUS_INACTIVE) && 
             (record->entries[i].status != MBR_STATUS_ACTIVE)) {
             cli_dbgmsg("cli_scanmbr: Invalid boot record status\n");
             return CL_EFORMAT;
         }
     }
 
     /* check the signature */
     if (record->signature != MBR_SIGNATURE) {
         cli_dbgmsg("cli_scanmbr: Invalid boot record signature\n");
         return CL_EFORMAT;
     }
 
     return CL_CLEAN;
 }
 
 /* this includes the overall bounds of extended partitions */
 static int mbr_primary_prtn_intxn(cli_ctx *ctx, struct mbr_boot_record mbr, size_t sectorsize)
 {
     prtn_intxn_list_t prtncheck;

     unsigned i = 0, pitxn = 0, prtncount = 0;

     int ret = CL_CLEAN, tmp = CL_CLEAN;

 
     prtn_intxn_list_init(&prtncheck);
 

     for (i = 0; i < MBR_MAX_PARTITION_ENTRIES && prtncount < ctx->engine->maxpartitions; ++i) {
         if (mbr.entries[i].type == MBR_EMPTY) {

             /* empty partition entry */

             prtncount++;
         }
         else {
             tmp = prtn_intxn_list_check(&prtncheck, &pitxn, mbr.entries[i].firstLBA,
                                         mbr.entries[i].numLBA);

             if (tmp != CL_CLEAN) {

                 if (tmp == CL_VIRUS) {

                     cli_dbgmsg("cli_scanmbr: detected intersection with partitions "
                                "[%u, %u]\n", pitxn, i);

                     ret = cli_append_virus(ctx, PRTN_INTXN_DETECTION);

                     if (SCAN_ALLMATCHES || ret == CL_CLEAN)

                         tmp = 0;
                     else
                         goto leave;
                 } else {

                     ret = tmp;

                     goto leave;

                 }
             }
 

             if (mbr.entries[i].type == MBR_EXTENDED) {
                 /* check the logical partitions */
                 tmp = mbr_extended_prtn_intxn(ctx, &prtncount, 
                                   mbr.entries[i].firstLBA, sectorsize);

                 if (tmp != CL_CLEAN) {

                     if (SCAN_ALLMATCHES && (tmp == CL_VIRUS)) {

                         ret = tmp;
                         tmp = 0;
                     }
                     else if (tmp == CL_VIRUS) {
                         prtn_intxn_list_free(&prtncheck);
                         return CL_VIRUS;
                     }
                     else {
                         prtn_intxn_list_free(&prtncheck);
                         return tmp;
                     }
                 }

             }
             else {
                 prtncount++;
             }
         }
     }

 leave:

     prtn_intxn_list_free(&prtncheck);
     return ret;
 }
 
 /* checks internal logical partitions */
 static int mbr_extended_prtn_intxn(cli_ctx *ctx, unsigned *prtncount, off_t extlba, size_t sectorsize)
 {
     struct mbr_boot_record ebr;
     prtn_intxn_list_t prtncheck;
     unsigned i, pitxn;

     int ret = CL_CLEAN, tmp = CL_CLEAN, mbr_base = 0;

     off_t pos = 0, logiclba = 0;

     int virus_found = 0;

     mbr_base = sectorsize - sizeof(struct mbr_boot_record);
 

     prtn_intxn_list_init(&prtncheck);

     logiclba = 0; i = 0;
     do {
         pos = extlba * sectorsize; /* start of extended partition */

         /* read the extended boot record */

         pos += (logiclba * sectorsize) + mbr_base;

         if (fmap_readn(*ctx->fmap, &ebr, pos, sizeof(ebr)) != sizeof(ebr)) {
             cli_dbgmsg("cli_scanebr: Invalid extended boot record\n");
             prtn_intxn_list_free(&prtncheck);
             return CL_EFORMAT;
         }

         /* convert the little endian to host */
         mbr_convert_to_host(&ebr);
 
         /* update state */
         (*prtncount)++;
 
         /* assume that logical record is first and extended is second */
         tmp = prtn_intxn_list_check(&prtncheck, &pitxn, logiclba, ebr.entries[0].numLBA);
         if (tmp != CL_CLEAN) {

             if (tmp == CL_VIRUS) {

                 cli_dbgmsg("cli_scanebr: detected intersection with partitions "
                            "[%u, %u]\n", pitxn, i);

                 ret = cli_append_virus(ctx, PRTN_INTXN_DETECTION);
                 if (ret == CL_VIRUS)
                     virus_found = 1;

                 if (SCAN_ALLMATCHES || ret == CL_CLEAN)

                     tmp = 0;
                 else
                     goto leave;
             } else {

                 ret = tmp;

                 goto leave;

             }
         }

         /* assume extended is second entry */
         if (ebr.entries[1].type != MBR_EXTENDED) {
             cli_dbgmsg("cli_scanebr: second entry for EBR is not an extended partition\n");
             break;

         }
 

         logiclba = ebr.entries[1].firstLBA;

         ++i;
     } while (logiclba != 0 && (*prtncount) < ctx->engine->maxpartitions);

  leave:
     prtn_intxn_list_free(&prtncheck);
     if (virus_found)
         return CL_VIRUS;

     return ret;

 }