#include <bastard.h>
#include <glob.h>

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


/* ---------------------------------------------------------------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.
    */
      target.header = malloc(strlen(strHeader)+ 1);
      strcpy(target.header, strHeader);
      disasm_env.flags |= DB_MOD;
      return(1);
   }

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

    /* +++++++++++++++++++ LEGACY LIBOPCODES STUFF ++++++++++++++++++ */
   int GetOpType( char *op){
   /* check operand type based on AT&T syntax, then strip AT&T syntax */
      char tmpStr[32];
   
      if ( op[0] == '%' ) { 
         return(OP_REG); }
      if ( op[0] == '$' ) { 
         return(OP_IMM); }
      if ( strchr(op, '(')) { 
         return(OP_IND); }
      if ( strncmp( op, "BYTE PTR", 8)) { 
         return(OP_BPTR); }
      if ( strncmp( op, "DWORD PTR", 9)) { 
         return(OP_DPTR); }
      if ( strncmp( op, "WORD PTR", 8)) { 
         return(OP_WPTR); }
      return(OP_UNK);
   }

	#ifndef NO_LIBOPCODES
   int disasmCB(FILE *stream, const char *format, ...){
   /* Replacement fprintf() for libopcodes.
    * NOTE: the following assumes src, dest order from disassembler */
      char *tmpStr;
      va_list args;
   
      va_start (args, format);
   
   /* this sucks, libopcodes passes one mnem/operand per call */
      if ( curr_inst->mnemonic[0] == 0) {
         strncpy(curr_inst->mnemonic, va_arg(args, char*), 15);
      } 
      else if ( curr_inst->src[0] == 0) {
         strncpy(curr_inst->src, va_arg(args, char*), 31);
         curr_inst->srcType = GetOpType(curr_inst->src);
      } 
      else if ( curr_inst->dest[0] == 0) {
         strncpy(curr_inst->dest, va_arg(args, char*), 31);
         curr_inst->destType = GetOpType(curr_inst->dest);
         if (strncmp(curr_inst->dest, "DN", 2) == 0)
            memset(&curr_inst->dest, 0, 32);
      } 
      else {
         strncpy(curr_inst->aux, va_arg(args, char*), 31);
         curr_inst->auxType = GetOpType(curr_inst->aux);
      }
      va_end (args);
      return(1);
   }

   int LibopWrap( BYTE * loc, BYTE table, struct code *inst, DWORD rva ){
   /* brought to you by the Libopcodes Blows Chunks Foundation */
      struct section *s = GetSectionObject(".text");
      target_info.buffer = target.image;
      target_info.buffer_length = target.size;
      target_info.buffer_vma = s->rva;
      curr_inst = inst;
      free(s);
      return((*disassemble_fn) (rva, &target_info));
   };
    /* +++++++++++++++++ END LEGACY LIBOPCODES STUFF ++++++++++++++++ */
	#endif

   int InitDisassembler(char *disasm){
   /*  This loads the disassembler of choice and calls its init routine
    *  Note that the use of libopcodes is still allowed; also everything
    *  is now force to dynamic-linking/loading */
      char libname[PATH_MAX];
      struct stat tmpStat;
   
   /* If a disassembler is loaded, unload it */
      if ( disasm_eng.init == TRUE && disasm_eng.lib ) {
         dlclose(disasm_eng.lib);
         disasm_eng.init = FALSE;
      }
   
   /* default to the preset disassembler */
      if ( disasm == NULL) disasm = disasm_eng.name;
		#ifndef NO_LIBOPCODES
      if ( strcmp(disasm, "libopcodes.so") == 0) {
      /* Allow libopcodes to be selected on the commandline or in a prefs
       * file, just in case the user is feeling self-destructive. For
       * now, the values will be hard-coded to Intel, but they should
       * be determined using the bfd stuff and the target struct settings.
       * NOTE: linking to BFD lets us use their assembler and such until
       * we make our own... */
      /* TODO : fix this ; it prob needs libiberty too */
         if ( ! (disasm_eng.lib = dlopen(disasm_eng.name, RTLD_NOW))  ){
            ErrorMsg(4200, dlerror());
            return(SetLastError(4200));
         } /* no libopcodes == lame distro */
         disassemble_fn = dlsym(disasm_eng.lib, "print_insn_i386_intel");
         disasm_eng.disasm_init = NULL;		/* no init routine needed */
         disasm_eng.disasm_addr = LibopWrap;	/* wrapper for libopcodes func */
      
      /* the INIT_DISASSEMBLER macro in dis-asm.h does the following: */
         target_info.flavour = bfd_target_unknown_flavour;
         target_info.arch = bfd_arch_i386;
         target_info.mach = 0;
         target_info.endian = BFD_ENDIAN_UNKNOWN;
         target_info.fprintf_func = disasmCB;
         target_info.stream = stdout;
         target_info.symbols = NULL;
         target_info.num_symbols = 0;
         target_info.buffer = target.image;
         target_info.buffer_length = target.size;
         target_info.buffer_vma = 0;
      //target_info.read_memory_func = buffer_read_memory;
      //target_info.memory_error_func = perror_memory;
      //target_info.print_address_func = generic_print_address;
      //target_info.symbol_at_address_func = generic_symbol_at_address;
         target_info.flags = 0;
         target_info.bytes_per_line = 0;
         target_info.bytes_per_chunk = 0;
         target_info.display_endian = BFD_ENDIAN_UNKNOWN;
         bfd_init();
      
         disasm_eng.init = TRUE;
         return(1);
      }	
		#endif
   
   /* OK, we are not using libopcodes, go ahead and load the library */
      sprintf(libname, "%s/lib%s.so", disasm_eng.path, disasm);
      if ( stat( libname, &tmpStat) != 0) {
      /* the user really hosed things up; force the i386 library */
         sprintf(libname, "%s/arch/libi386.so", disasm_env.home);
      }
      if ((disasm_eng.lib = dlopen(libname, RTLD_NOW)) == NULL){
         ErrorMsg(4200, dlerror());
         return(SetLastError(4200));
      }
		/* Load the two important routines */
      disasm_eng.disasm_init = dlsym(disasm_eng.lib, "DisasmInit");
      if ( disasm_eng.disasm_init)  (*disasm_eng.disasm_init)(NULL);
      disasm_eng.disasm_addr = dlsym(disasm_eng.lib, "DisasmAddress");
      if (! disasm_eng.disasm_addr)    {ErrorMsg(603, dlerror());
         return(-1);}
		/* load the rest of the disasm lib */
		disasm_eng.get_prologue = dlsym(disasm_eng.lib, "GetFunctionPrologue");
		disasm_eng.get_epilogue = dlsym(disasm_eng.lib, "GetFunctionEpilogue");
		disasm_eng.get_endian   = dlsym(disasm_eng.lib, "GetByteOrder");
		disasm_eng.get_addrsize = dlsym(disasm_eng.lib, "GetAddrSize");
		disasm_eng.get_byte     = dlsym(disasm_eng.lib, "GetByteSize");
		disasm_eng.get_word     = dlsym(disasm_eng.lib, "GetWordSize");
		disasm_eng.get_dword    = dlsym(disasm_eng.lib, "GetDwordSize");
		disasm_eng.get_SP       = dlsym(disasm_eng.lib, "GetSP");
		disasm_eng.get_IP       = dlsym(disasm_eng.lib, "GetIP");

      disasm_eng.init = TRUE;
      return(1);
   }

   int AnalyzeCurrOp( struct code *code, char *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;
      int addrType, xrefType;
      char addrName[64] = "\0";
      long rva = 0;
   
   /* Mask the XREF Type [rwx] from the instruction */
      xrefType = type & 0x07;
   
   /* TODO: string recognition? e.g. if OP_STRING_PTR on rep scasb? */
	/* TODO: change operands && operand types to absolute addresses */
      switch (type & 0x0F00 ) {
         case OP_REG:
         case OP_IND:
         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:
         /* RELATIVE OFFSET Convert offset to RVA */
            if ( a = GetAddressObject(code->rva) ) {
            /* add IMM value to end of current instruction [rva + size] */
               rva = strtol(op, (char **)NULL, 10) + code->rva + a->size;
               free(a);
            } 
            break;
         case OP_STRPTR: 
         /* for now, just fall through -- later, create string */
         case OP_BPTR: case OP_WPTR: case OP_DPTR:
         /* ABSOLUTE ADDRESS convert ASCII rep of RVA to binary */
            rva = strtol(op, (char **)NULL, 16);
            break;
         default:
         /* why are so many operands '0'? Check libi386 code. */
            break;
      }
		/* NOTE : if rva, move RVA into operand; move operand name, value
		 *        and type into a comment. Change operand type to OP_ADDR
		 */
   //if (IsValidAddr(rva)){
      if (rva){
      /* found a program address in instruction */
         addrType = ADDR_DATA;
      
         if (code->mnemType & INS_SUB){ 
         /* make target a function */ 
         /* NOTE do strack frames 7 'push ebp; mov ebp, esp' in pass */
            addrType = ADDR_CODE;
            if ( d_keyfind(FUNCTION_RVA, &rva) != S_OKAY){
               sprintf(addrName, "sub_%08X", rva);
               DefineFunc( rva, addrName, 0 );
            }
         } 
         else if (	code->mnemType & INS_BRANCH ||
         code->mnemType & INS_COND ) {
         /* This is an executable addr */
            addrType = ADDR_CODE;
         /* Do nothing :) */
         } 
      
      /* create an address object for it if none exists */
         if (d_keyfind(ADDRESS_RVA, &rva) != S_OKAY)
            DefineAddr( rva, 1, 0, addrType);
      /* If not already named, name it */
         if ( ! addrName[0] ) sprintf(addrName, "loc_%08X", rva);
         if(d_keyfind(NAME_RVA, &rva) != S_OKAY) 
            SetName( rva, addrName, NAME_AUTO); 
      /* Make an xref from this instruction to the rva */
         DefineRef(code->rva, rva, xrefType);
      }
      return(1);
   }

   int AnalyzeCurrInst(struct code *code){
   /* If disasm->GetComment(mnemonic) ...
   /* Get Operand Types */
      if (code->dest) AnalyzeCurrOp( code, code->dest, code->destType);
      if (code->src) AnalyzeCurrOp( code, code->src, code->srcType);
      if (code->aux) 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 DisassembleFull(){
      struct function f;
   /* disassemble forward from entry point */
      if (! disasm_eng.init) InitDisassembler( target.cpu );
   
   /* Start by disassembling all functions recognized by the format parser */
   /* TODO : This part is VERY SLOW -- figure out why */
      Message("Disassembling named symbols\n");
      d_keyfrst(FUNCTION_RVA);
      while (db_status == S_OKAY) {
         d_recread(&f);
         if (db_status == S_OKAY) DisassembleForward(f.rva);
         d_keynext(FUNCTION_RVA);
      }
   
   /* Then attack the entry point */
      Message("Disassembling forward from entry point\n");
      DisassembleForward(target.entry);
	/* create an export for the entry point */
		DefineExport(target.entry, "_start");
      Message("Performing post-disassembly passes\n");
      DisassembleAllPasses();
      Message("Done.\n");
      return(1);
   }

   int DisassembleDumb( ){
      struct section s;
   
      d_keyfrst(SECTION_RVA);
      while (db_status == S_OKAY) {
         d_recread(&s);
         if ( s.flags & SECTION_CODE || s.flags & SECTION_EXECUTE)
            DisassembleSection( s.name );
         d_keynext(SECTION_RVA);
      }
      return(1);
   }

   int DisAddr(DWORD pa, struct code *c, DWORD rva){
   /* pa  = offset from image start  
    * c   = code struct filled with info
    * rva = RVA to disassemble */
      struct address *a;
      int bytes, size;
   
      if ( d_keyfind(CODE_RVA, &c->rva) == S_OKAY){
      /* code object already exists; skip over it */
         if ( a = GetAddressObject(c->rva) ){
            size = a->size;
            free(a);
            return(size);
         } 
         else 
            return(SetLastError(4210));
      } 
      else {
         bytes = (*disasm_eng.disasm_addr)( target.image + pa, 0, c, rva) ; 
         if (! strcmp(c->mnemonic, "invalid") || bytes < 1) 
            return(-1);
      /* Make an address object */
         DefineAddr(c->rva, bytes, pa, ADDR_CODE);
      /* Make a code object */
         AnalyzeCurrInst(c);
      }
      return(bytes);
   }

  int DisassembleSection(char *name){
      struct section *sec;
      char msg[128];
   
      printf("enter disassemble section");
      if (! disasm_eng.init) InitDisassembler(NULL);
      if ( sec = GetSectionObject(name)) {
         sprintf(msg, "Disassembling section %s: %X\n", name, sec->rva);
         Message(msg);
         DisassembleRange(sec->rva, sec->size);
         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 (! disasm_eng.init) InitDisassembler(NULL);
      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;
      long new_rva;
		long tmp_rva;
   
      if (! disasm_eng.init) InitDisassembler(target.cpu);
      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)){
         //rva is within segment
            if ( d_keyfind(ADDRESS_RVA, &sec.rva) == S_NOTFOUND) {
               d_keynext(ADDRESS_RVA);
            }
            d_recread(&addr);
            addr.rva = rva;
            addr.pa += (rva - sec.rva);
            addr.flags = ADDR_CODE;
            while (cont){
               memset(&code, 0, sizeof(struct code));
               code.rva = rva + bytes;
               size = DisAddr(addr.pa + bytes, &code, code.rva);
               if (size < 1) cont = 0;
               else {
               //DefineAddr(curr_inst.rva, size, addr.pa + bytes, ADDR_CODE);
                  bytes += size;
                  switch( code.mnemType ) {
                     case INS_BRANCH :
                        cont = 0;
                     case INS_COND :
                     case INS_SUB :
                     //printf("Branch Type %d ", curr_inst.destType);
                        if (code.destType & OP_REG) 
                        /* skip registers for now */
                           break;
                        else if (code.destType & OP_IMM || 
													 code.destType & OP_REL)
								/* convert relative offset to address */
                           new_rva = strtol(code.dest, NULL, 16)+ size + code.rva;
                        else
                        /* convert char address to long */
                           new_rva = strtol(code.dest, NULL, 16);	
                     // printf("following branch %s %X\n",code_inst.dest,new_rva);
                        DisassembleForward( new_rva );
                        break;
                     case INS_RET :
                        cont = 0;
                     default:
                        break;
                  }	// end switch
               }		// end if (size)
               if ( addr.rva + bytes >= sec.rva + sec.size) cont = 0;
            } 			// end while (cont)
         }				// end if (rva)
         d_keynext(SECTION_NAME);
      }					// end while (db_status)
      disasm_env.flags |= DB_MOD;
      return(1);
   }

   int DisassemblePass( int pass ){
      char tmp[PATH_MAX];
   
      sprintf( tmp, "%sengines/pass%d.sc", disasm_env.home, pass);
      return(RunScript(tmp));
   }

   int DisassembleAllPasses(){
      char tmp[4096], *pos;
      glob_t globbuf;
      int x, len;
   
      sprintf(tmp, "%sengines/pass*.sc", 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],".sc") - pos;
         strncpy(tmp, pos, len);
         tmp[len] = '\0';
         DisassemblePass( atoi(tmp) );
      }
      globfree(&globbuf);
      return(1);
   }


	/* Convenience functions -- low-level wrappers to the disasm_lib fn's */
int DisPrologue(struct code **c) {
	if (! disasm_eng.get_prologue) return(0);
	return((*disasm_eng.get_prologue)(c));
}
int DisEpilogue(struct code **c) {
	if (! disasm_eng.get_epilogue) return(0);
	return((*disasm_eng.get_epilogue)(c));
}
int DisByteOrder() {
	if (! disasm_eng.get_endian) return(0);
	return((*disasm_eng.get_endian)());
}
int DisAddrSize() {
	if (! disasm_eng.get_addrsize) return(0);
	return((*disasm_eng.get_addrsize)());
}
int DisByteSize() {
	if (! disasm_eng.get_byte) return(0);
	return((*disasm_eng.get_byte)());
}
int DisWordSize() {
	if (! disasm_eng.get_word) return(0);
	return((*disasm_eng.get_word)());
}
int DisDwordSize() {
	if (! disasm_eng.get_dword) return(0);
	return((*disasm_eng.get_dword)());
}
char * DisGetSP() {
	if (! disasm_eng.get_SP) return(0);
	return((*disasm_eng.get_SP)());
}
char * DisGetIP() {
	if (! disasm_eng.get_IP) return(0);
	return((*disasm_eng.get_IP)());
}
