#include <api/api_main.h>
#include <extension.h>
#include <debug.h>

extern struct DISASM_TGT target;
extern struct DISASM_ENV disasm_env;
extern struct DISASM_PREFS disasm_prefs;

/* API sections:
 *     4. Disassembly         :: All disassembly passes
 */


struct code *curr_inst;
/* ---------------------------------------------------------------Disassembly */
int CreateHeader( char *strHeader ){
/* This just stores a character representation of the file header 
 * in memory -- formatting is left up to the file format parser
 * script.
 */
   if (target.header) free (target.header);
   target.header = malloc(strlen(strHeader)+ 1);
   if (! target.header ) return(SetLastError(9050));
   strcpy(target.header, strHeader);
   disasm_env.flags |= DB_MOD;
   return(1);
}

void PrintHeader() {
   if (target.header) printf("%s\n", target.header);
   return;
}

int deal_with_addr( unsigned long rva, struct code *code, int xrefType, 
                    char *prefix) {
   int addrType; 
   char addrName[32];

   if ( rva && IsValidAddr(rva)){
      addrType = ADDR_DATA;

      if (code->mnemType & INS_SUB){ 
         /* make target a function */ 
         addrType = ADDR_CODE;
         DefineFunc( rva, NULL, 0, 0);
      } else if ( code->mnemType & INS_BRANCH ||
                  code->mnemType & INS_COND ||
                  xrefType & XREF_EXEC) {
         /* This is an executable addr */
         addrType = ADDR_CODE;
      }

      /* create an address object for it if none exists */
      if (d_keyfind(ADDRESS_RVA, &rva) != S_OKAY)
         DefineAddr( rva, 4, 0, addrType);

      /* If not already named, name it */
      if (d_keyfind(NAME_RVA, &rva) != S_OKAY) {
         sprintf(addrName, "%s_%08X", prefix, rva); 
         SetName( rva, addrName, NAME_AUTO); 
      }

      /* Make an xref from this instruction to the rva */
      DefineRef(code->rva, rva, xrefType);
   }
   return(1);
}

int AnalyzeCurrOp( struct code *code, long *op, int *type){
/* This finds absolute and rfelative address references in an opcode,
 * creates an address object for that reference, creates an xref to
 * that address object, and performs location naming/subroutine
 * definition for references that are jmp or call targets */
   struct address *a;
   struct comment *cmt;
   int xrefType, change = 0;
   char buf[64] = "\0";
   long rva = 0;

   /* Mask the XREF Type [rwx] from the instruction */
   xrefType = *type & OP_PERM_MASK;

   /* first do a check and be sure this is not a valid offset disguised as
    * an IMM/REL [e.g. 'pop [addr]' is considered an IMM] */
   if ( ! (*type & OP_PTR) && IsValidCodeAddr(*op)){
      *type &= ~OP_TYPE_MASK;         /* 0xFF0FF0FF;  clear Type */
      *type |= OP_ADDR;  
   }

   switch (*type & OP_TYPE_MASK ) {
      case OP_REG:
            /* Define a Register Effect if the op is a REG and is Written */
            if (xrefType == OP_W) {
               /* if we know the change amount, use it -- else 0 */
               if ( code->src && (code->srcType & OP_TYPE_MASK == OP_IMM ))
                  change = code->src;
               DefineCodeEffect(code->rva, *op, change);
            }
         break;
      case OP_EXPR:
         /* for now, ignore registers && indirect addresses */
         break;
      case OP_IMM: 
         /* If IMM is only read, it is a constant */
         if ( !(xrefType & XREF_EXEC) && !(xrefType & XREF_WRITE) ) 
            break;
         /* else fall through to 'rel offset' */
      case OP_REL:
/* TODO: For DATA, associate default data type with target rva */
         /* RELATIVE OFFSET Convert offset to RVA */
         if ( a = GetAddressObject(code->rva) ) {
            /* add IMM value to end of current instruction [rva + size] */
            rva = *op + code->rva + a->size;

            /* note the change in instruction comment */
/* TODO: why does push eax end up with this coment? */
            sprintf(buf, "(0x%X was %+d)", rva, *op);
            SetAddrComment(code->rva, DefineComment(buf, CMT_AUTO));

            /* change the operand to the new rva */
            *op = rva;
            *type &= ~OP_TYPE_MASK;         /* 0xFF0FF0FF;  clear Type */
            *type |= OP_ADDR;                /* Set type to Absolute Addr */
            free(a);

         }  else {
            /* This shouldn't happen */
            DEBUG_PrintMsg("Horrible error in AnalyzeCurrOp! "
                           "Addr %x for CODE struct doesn't exist\n", 
                           code->rva); ;
         break;
         }
         /* fall through to handle new address */
      case OP_ADDR: 
/* TODO: For DATA, associate default data type with rva */
         rva = *op;
         deal_with_addr( rva, code, xrefType, "loc");
         break;
      case OP_PTR: 
/* TODO: For DATA, associate default pointer type with pointer rva */
/* TODO: For DATA, associate default data type with deferenced rva */
         rva = *op;
         /* first, create an address object for the pointer itself */
         deal_with_addr( rva, code, xrefType, "ptr");

         /* next, create an address for the rva contained in the pointer */

         /* find section that rva is in, get offset from start of section,
          * read bytes from file, use *that* as rva.Put op in brackets. */
         if (rva = GetPA(rva) && rva < target.size){
            /* get address in pointer from target image, if possible */
            encpy((char *)&rva, target.image + rva, ext_arch->sz_addr);
            if (IsValidAddr(rva))
               DefineRef(code->rva, rva, xrefType);
         }
         break;
      default:
         break;
   }

   if ((*type & OP_MOD_MASK ) &  OP_STRING){ 
/* TODO: string recognition? e.g. if OP_STRING_PTR on rep scasb? */
      rva = *op;
      DefineString(rva);
   }

   return(1);
}

int AnalyzeCurrInst(struct code *code){
   struct code_effect e;
/* If disasm->GetComment(mnemonic) ... */
/* Get standard reg effects */
   if ( ext_gen_effect(code->mnemonic, &e)) {
      DefineCodeEffect(code->rva, e.reg, e.change);
   }
/* Get Operand Types */
   if (code->destType) AnalyzeCurrOp( code, &code->dest, &(code->destType) );
   if (code->srcType) AnalyzeCurrOp( code, &code->src, &(code->srcType) );
   if (code->auxType) AnalyzeCurrOp( code, &code->aux, &(code->auxType) );

/* Get Instruction Type */

/* Make Code object */
   if ( d_fillnew(CODE, code) != S_OKAY){
      return(SetLastError(dbError()));
   }

   return(1);
}

int DisassembleTarget(char *disasm, void *param){
   /* called with libdisasm.$DISASM.so as 'disasm',
    * e.g. "libdisasm.full.so would be "full" */
   char fullname[256] = {0};

   
   DEBUG_PrintMsg("Disassemble %s on TARGET ", disasm);
   DEBUG_PrintMsg("entry %08X\n", target.entry);
   sprintf(fullname, "disasm.%s", disasm);
   if ( LoadPlugin(fullname)) 
      ext_plugin_main( param );
/* else error */
   return(1);
}
int DisAddr(long pa, struct code *c, long rva){
/* pa  = offset from image start  
 * c   = code struct filled with info
 * rva = RVA to disassemble */
   struct address *a;
   struct code code;
   int bytes;


   /* IF CODE object for/containing rva already exists 
   if ( DBTableFirst(CODE_RVA, &code) &&           /* check table !empty 
         ( DBIndexFind(CODE_RVA, &rva, &code) ||   /* rva a CODE object? 
            ( DBIndexPrev(CODE_RVA, &code )  && 
              DBIndexFind(ADDRESS_RVA, &code.rva, &a) &&
              code.rva <= rva && code.rva + a.size > rva)  )  ) 
   { /* rva in CODE object? 

      if ( disasm_prefs.options & DISASM_REDO ) {  
         /* Undefine Code Object for a->rva 
         UndefineCode(code.rva);
      } else {
         if (code.rva == rva) {
            DBIndexFind(ADDRESS_RVA, &code.rva, &a);
            memcpy(c, &code, sizeof(struct code));
            return(a.size);
         } else return (  a.size - (rva - a.rva) );
      }
   } */
   if (d_keyfind(CODE_RVA, &rva) == S_OKAY) {
      if (disasm_prefs.options & DISASM_REDO) {
         UndefineCode(rva);
      } else if (a = GetAddressObject(rva)){
         bytes = a->size;
         free(a);
         //return(bytes);
         return(0);
      } else
         return(SetLastError(4210));
   }
         

   /* IF this is an invalid Physical Address, BailOut */
   if (  pa < 0 || pa > (long)target.size) return(-1);
   /* ELSE call disassembler (ARCH)  extension */
   bytes = ext_disasm_addr(target.image + pa, 0, c, rva); 

   /* IF this is an invalid instruction, BailOut */
   if ( bytes < 1) 
         return(-1);

   /* ELSE Make an address object for rva */
   c->rva = rva;
   DefineAddr(c->rva, bytes, pa, ADDR_CODE); /* Define handles dupes */

   /* Analyze current instr and create a CODE object */
   AnalyzeCurrInst(c);
   
   return(bytes);    /* return size of instruction */
}

  int DisassembleSection(char *name){
   struct section *sec;
   struct address a;
   char msg[128];

   if (ext_arch->ext.flags & EXT_STATUS_MASK != EXT_INIT) 
      SetTargetArch( target.arch );
   if ( sec = GetSectionObject(name)) {
      sprintf(msg, "Disassembling section %s: %X\n", name, sec->rva);
      Message(msg);
      if ( DBFindClosestNext(ADDRESS_RVA, &sec->rva, &a) )
         DisassembleRange(a.rva, sec->size - (a.rva - sec->rva));
      Message("Done.\n");
      free(sec);
   } 
   else  
      return(0); /* error will be set by GetSectionObject */
   disasm_env.flags |= DB_MOD;
   return(1);
}

int DisassembleRange(long rva, int size){
/* rva must have an address record */
   struct address *addr;
   struct code code;
   int bytes, len;

   if (ext_arch->ext.flags & EXT_STATUS_MASK != EXT_INIT) 
      SetTargetArch( target.arch );
   disasm_env.flags |= DB_MOD;
   if (! (addr = GetAddressObject(rva)) ) 
      return(0); /* last_err already set */
   for (bytes = 0; bytes <= size; ){
      memset(&code, 0, sizeof( struct code));
      code.rva = rva + bytes;
      len = DisAddr(addr->pa + bytes, &code, code.rva);
      if (len < 1) { 
         bytes++; 
         break;
      } else bytes += len;
   }
   free(addr);
   return(1);
}

int DisassembleForward(long rva){
   struct section sec;
   struct address addr;
   struct code code;
   int bytes=0, size=0, cont = 1, destType, *state;
   char tmp_op[32]; /* temporary operand */
   long new_rva, tmp_rva;

   if (ext_arch->ext.flags & EXT_STATUS_MASK != EXT_INIT) 
      SetTargetArch( target.arch );
   state = DBSaveState();
   d_keyfrst(SECTION_NAME);
   while ( db_status == S_OKAY ) {   
      d_recread(&sec);   
      tmp_rva = sec.rva + sec.size;
      if (db_status == S_OKAY && 
          rva >= sec.rva && rva < (sec.rva + sec.size)){
         /* Get starting addr of segment to use as a base for pa */
         if ( d_keyfind(ADDRESS_RVA, &rva) == S_NOTFOUND) {
            d_keyprev(ADDRESS_RVA);
         }
         d_recread(&addr);
         addr.pa += (rva - addr.rva);
         addr.rva = rva;
         addr.flags = ADDR_CODE;
         Message(" \b+"); /* Display a visual stack */
         while (cont){
            memset(&code, 0, sizeof(struct code));
            code.rva = rva + bytes;
            Message("\b|"); /* step 1 of this stack level */
            size = DisAddr(addr.pa + bytes, &code, code.rva);
            Message("\b/"); /* step 2 on this stack level */
            if (size < 1) cont = 0;
            else {
            /* At this point, DefineAddr and AnalyzeInst have been done */
               new_rva = 0;
               bytes += size;
               destType = code.destType & 0x0F00;
               switch( code.mnemType ) {
                  case INS_BRANCH :
                     cont = 0;
                  case INS_COND :
                  case INS_SUB :
                  //printf("Branch Type %d ", code.Type);
                     if (destType ==  OP_REG ) {
                     /* skip registers for now */
                        ;
                     } else if (destType == OP_PTR) {
                        new_rva = code.dest;
                        /* 1. get pa for rva */
                        if ( new_rva = GetPA(new_rva) ) {
                        /* 2. get dword [pointer] stored at pa */
                        /* this is cpu-specific and needs to be changed: */
                           encpy((char *)&new_rva, target.image + 
                                       new_rva, ext_arch->sz_addr);
                        } 
                     } else if (destType == OP_IMM || destType == OP_REL) {
                        /* convert relative offset to address -*/
                        new_rva = code.rva + size + code.dest;
                     } else {
                        new_rva = code.dest;   
                     }
                     break;
                  case INS_RET :
                     cont = 0;
                  default:
                     break;
               }   // end switch
               /* Follow branch if one was detected */
               Message("\b-"); /* step 3 on this stack level */
               /* put a stack bounds check in here ??? */
               /* TODO : IsValidAddr and such? */
               if (new_rva && new_rva != rva) {
                  DisassembleForward( new_rva );
               }
            }      // end if (size)
            Message("\b/"); /* step 4 on this stack level */ 
            if ( addr.rva + bytes >= sec.rva + sec.size) 
               cont = 0;
         }          // end while (cont)
         Message("\b \b"); /* undo `Message("|")` */
      }            // end if (rva)
      d_keynext(SECTION_NAME);
   }               // end while (db_status)
   disasm_env.flags |= DB_MOD;
   DBRestoreState(state);
   return(1);
}

int DisassemblePass( int pass , int options){
   char tmp[PATH_MAX];

   sprintf( tmp, "%s/engines/pass%d.bc", disasm_env.home, pass);
   if (! LoadExtension(EXT_ENGINE, tmp, ext_engine))
      return(SetLastError(4593));
   return(ext_engine_main(&options));
}

int DisassembleAllPasses(){
   char tmp[4096], *pos;
   glob_t globbuf;
   int x, len;

   sprintf(tmp, "%s/engines/pass*.bc", disasm_env.home);
   glob(tmp, 0, NULL, &globbuf);
   for (x=0; x < globbuf.gl_pathc; x++){
      pos = strstr(globbuf.gl_pathv[x],"/pass") + 5;
      len = strstr(globbuf.gl_pathv[x],".bc") - pos;
      strncpy(tmp, pos, len);
      tmp[len] = '\0';
      DisassemblePass( atoi(tmp), 0 );
   }
   globfree(&globbuf);
   return(1);
}


/* Convenience functions -- low-level wrappers to the disasm_lib fn's */
int DisPrologue(struct code **c) {
return(ext_get_prologue(c));
}
int DisEpilogue(struct code **c) {
return(ext_get_epilogue(c));
}
int DisByteOrder() {
return(ext_arch->endian);
}
int DisAddrSize() {
return(ext_arch->sz_addr);
}
int DisByteSize() {
return(ext_arch->sz_byte);
}
int DisWordSize() {
return(ext_arch->sz_word);
}
int DisDwordSize() {
return(ext_arch->sz_dword);
}
int DisGetSP() {
return(ext_arch->SP);
}
int DisGetIP() {
return(ext_arch->IP);
}
