#include <bastard.h>
#include <readline/readline.h>
#include <seer.h>
#include <sys/types.h>


/* API sections:
 *     1. Loading 				:: Loading files, DBs, and so on
 *     2. Saving					:: Saving DBs, asm listings, patched files
 * /

/* ---------------------------------------------------------------Loading */
   int LoadTarget(char *targetName ) {
   /* MMapping might cause problems with larger files; however it 
    * should have smaller storage requirements than actually loading
    * the bytes into the database itself. It might be better to mmap each
    * section, thus junk like .bss can be summarized instead of stored.
    */
      char *tmpStr, cmdBuf[128];
      FILE *fTemp;
      DIR *dirBDB;
      struct stat tmpStat;
   
   //  Make sure no DB is currently loaded
      if (disasm_env.flags & DB_LOADED) CloseDB();
   
   //   1. make .bdb temp dir for target database
      if ( mkdir(".bdb", 0777) == -1 ){
         return(SetLastError(4510));	
      }
      if ( (dirBDB = opendir(".bdb")) == NULL ) {
         return(SetLastError(4520));	
      } 
      else closedir(dirBDB);
   
   //   2. make target_path and target
      tmpStr = strrchr(targetName, '/');
      if (tmpStr != NULL) {
         strncpy(target.path, targetName, tmpStr - targetName + 1);
         strncpy(target.name, (tmpStr + 1), sizeof(target.name));
         target.path[tmpStr - targetName + 1] = 0; //terminate string
      } 
      else {
         strncpy(target.path, "./", 3);	
         strncpy(target.name, targetName, sizeof(target.name));	
      }
   
   //   3. Copy target binary to .bdb
      sprintf(target.dbname, "./%s.bdb", target.name);
      sprintf(cmdBuf, "cp %sbdb.dbd ./.bdb/%s.dbd", disasm_env.home, 
         target.name);
      if ( pclose( popen( cmdBuf, "r"))  == -1){
         return(SetLastError(4530));
      }
      sprintf(cmdBuf, "cp -f %s%s ./.bdb/%s", target.path, 
         target.name, target.name);
      if ( pclose( popen( cmdBuf, "r"))  == -1){
         return(SetLastError(4540));
      }
   
   //   4. memory map target image
      sprintf(cmdBuf, "./.bdb/%s", target.name);
      target.FD = open(cmdBuf, O_RDWR);
      if (target.FD == -1) { 
         sprintf(cmdBuf, "Cannot open file %s", errno);
         ErrorMsg(9030, cmdBuf);
         return(SetLastError(9030));
      }
      disasm_env.flags |= TGT_MAPPED;
      fstat(target.FD, &tmpStat);
      target.size = tmpStat.st_size;
      target.image = mmap(0, target.size, PROT_READ | PROT_WRITE, 
         MAP_SHARED, target.FD, 0);
   
   //   5. Load Database
      printf("loading %s\n", target.name);
      if( d_open(target.name, "s") != S_OKAY ) {
         sprintf(cmdBuf, "DB error: Cannot open database (db_status %d)\n", 
            db_status);
         ErrorMsg( 4550, cmdBuf);
         return(SetLastError(4550));
      } 
      disasm_env.flags |= DB_LOADED || DB_MOD;
      initSequences();
      return(1);
   }

   int LoadBDB(char *dbname) {
   /* Load a previously-saved .bdb 
    */
      struct stat s;
      FILE *f;
      char cmdBuf[256];
   
   //  Make sure no DB is currently loaded 
      if (disasm_env.flags & DB_LOADED) CloseDB();
   
   //  -- if dbname is NULL then use target.name
   // 1. if exists $target.bdb exec tar -zxf $target.bdb -C./bdb
      if ( ! stat(dbname, &s) ) {
         sprintf(cmdBuf, "tar -zxf %s", dbname);
         if ( pclose( popen( cmdBuf, "r"))  == -1){
            return(SetLastError(4560));
         }
      } 
      else 
         return(SetLastError(9020));
   
   // 2. direct Typhoon to that DB
      if( d_open(target.name, "s") != S_OKAY ){
         sprintf(cmdBuf, "DB error: Cannot open database (db_status %d)\n", 
            db_status);
         ErrorMsg( 4550, cmdBuf);
         return(SetLastError(4550));
      }
   // 3. Load file header
      if ( ! stat("./.bdb/header.txt", &s)) {
         f = fopen( "./.bdb/header.txt", "r");
         if (f) {
            target.header = calloc(s.st_size, 1);
            fread( target.header, s.st_size, 1, f);
            fclose(f);
         }
      }
   // 4. Fill Target Struct
      if (! target.header ) {
         target.header = calloc( 64, 1);
         sprintf(target.header, "Header not available!\n");
      }
   //target.size, target.image
      initSequences();
      disasm_env.flags |= DB_LOADED;
      return(1);
   }

   int SetFileFormat(char* file_format){
   /* Invoke the parser in $BASTARD/formats
    * This parser should parse the target file and create sections for
    * all data and code sections/segments in the file. Additional info,
    * such as relocations/imports/exports, should be dealt with by the
    * parser as well.
    */
      char ff_macro[128], errBuf[256];
      struct stat statbuf;
      int ff;
      scScript ff_script;
      scInstance *ffInst;
   
      sprintf(ff_macro, "%sformats/%s.sc", disasm_env.home, file_format);
      if (stat( ff_macro, &statbuf) == -1) {
         sprintf(errBuf, "Unsupported File Format %s!\n", ff_macro);
         ErrorMsg( 4570, errBuf);
         return(SetLastError(4570));
      }
		RunScript(ff_macro);

      disasm_env.flags |= DB_MOD;
      return(1);
   }

   int SetFileArch(char *file_arch) {
      if (strcmp( disasm_eng.name, "libopcodes.so") )
         strncpy(disasm_eng.name, file_arch, 31); 
      InitDisassembler( NULL );
      disasm_env.flags |= DB_MOD;
      return(1);
   }
/* ---------------------------------------------------------------Saving */
   void SaveDBAs( char *name) {
   /* store database and image in .bdb file */
      char cmdBuf[256];
      FILE *f;
   /* create .info file */
   /* create header str file */
      f = fopen( "./.bdb/header.txt", "w");
      if (f) {
         fprintf( f, "%s", target.header);
         fclose(f);
      }
      msync(0, target.size, MS_SYNC);
      sprintf(cmdBuf, "tar --remove-files -zcf %s ./.bdb", name);
      if ( pclose( popen( cmdBuf, "r"))  == -1) {
         ErrorMsg(4565, "Unable to save target .bdb!");
         SetLastError(4565); 
         return;
      }	
      else Message("DB saved\n");	
   }

   void SaveDBBak() {
      char fname[256];
      sprintf(fname, "./%s.%d.bdb", target.name, time(NULL));
      SaveDBAs( fname );
      return;
   }

   void SaveDB( ) {
      char *str;
      int c;
   
      if ( disasm_env.flags & ANNOY_USER ) {
         printf("Save %s database? <Y/n/a/b/?> ", target.name);
         c = getc( stdin ); getc(stdin); 			//strip CR
      
         switch (c) {
            case 'n': case 'N':
               return;
            case 'b': case 'B':
               SaveDBBak(); 
               return;
            case 'a': case 'A':
               str = readline("Save as: ");
               strncpy(target.dbname, str, 255);
               free(str);
               break;
            case 'y': case 'Y': case '\n':
               break;
            case '?': case 'h': default:
               printf("  Y='Yes' N='No' A='save As...' B='make Backup'\n");
               SaveDB();
               return;
         }
      }
      SaveDBAs(target.dbname);
      disasm_env.flags ^= DB_MOD;
      return;
   }

   void CloseDB( ) {
      char cmdBuf[128];
   
      if (disasm_env.flags & DB_MOD) {			
         SaveDB();
         disasm_env.flags ^= DB_MOD;
      }
      if (disasm_env.flags & DB_LOADED) {
         d_close();										/* close Typhoon DB */
         disasm_env.flags ^= DB_LOADED;
      }
      if (disasm_env.flags & TGT_MAPPED) {
         close(target.FD);								/* cleanup mmap'ed image */
         munmap(target.image, target.size);
         disasm_env.flags ^= TGT_MAPPED;
      }
      if (target.header) free(target.header);
      memset( &target, 0, sizeof(target));
      memset( &disasm_eng, 0, sizeof(disasm_eng));
      sprintf(cmdBuf, "rm -rf ./.bdb");
      if ( pclose( popen( cmdBuf, "r"))  == -1) 
         ErrorMsg(4511, "Error cleaning up .bdb");
   /* restore defaults */
      strcpy(target.name, "a.out");
      strcpy(target.cpu, "i386");
      strcpy(target.format, "ELF");
      strcpy(disasm_eng.name, "libi386.so");
      sprintf(disasm_eng.path, "%s/arch", disasm_env.home);
   
      return;
   }

   int SaveAsASM(char *filename){
      FILE *asmFile;
      char *line;
      struct address addr;
      struct section sec;
   
      if (! filename ) 
         return(SetLastError(9010));
      asmFile = fopen(filename, "w");
      if ( ! asmFile) 
         return(SetLastError(9030));
      line = calloc(1024, 1);
   
   /* print target info */
      fprintf( asmFile, "Disassembly of %s by the bastard\n\n", target.name);
   /* print header */
      fprintf( asmFile, "%s", target.header);
   /* print each section */
      d_keyfrst(SECTION_RVA);
      while( db_status == S_OKAY) {
         d_recread(&sec);
         fprintf(asmFile, "Section: %s\tStart Address: %08X\tEnd Address: %08X\n",
            sec.name, sec.rva, sec.rva + sec.size);
         fprintf(asmFile, "-------------------------------------------------"
            "------------------------\n");
         d_keyfind(ADDRESS_RVA, &sec.rva);
         while( db_status == S_OKAY && d_recread(&addr) == S_OKAY &&
         addr.rva < (sec.rva + sec.size) ){
            if (addr.flags & ADDR_CODE)
               asmsprintf(line, disasm_env.asmFmt, &addr);
            //else datasprintf(line, NULL, &addr);
            else sprintf(line, "%08X\t Data region of size %d", 
                  addr.rva, addr.size);
            fprintf(asmFile, "%s\n", line);
            d_keynext(ADDRESS_RVA);
         }
         fprintf(asmFile, "\n");
         d_keynext(SECTION_RVA);
      }
      fclose( asmFile );
      free(line);
      return(1);
   }

   int SaveAsDump(char *filename){
      FILE *dumpFile;
      unsigned int x, y;
      unsigned char hex[48+4], ascii[16+4], c;
   
      if (! filename ) 
         return(SetLastError(9010));
      dumpFile = fopen(filename, "w");
      if (! dumpFile ) 
         return(SetLastError(9030));
   
      for (x = 0; x < target.size; x += y) {
         memset(hex, 0, 52);
         memset(ascii, 0, 20);
         for (y = 0; x + y < target.size && y < 16; y++) {
            c = ((char *)target.image)[x + y];
            sprintf(hex, "%s%02X ", hex, c); 
            if ( isprint(c) ) ascii[y] = c;
            else ascii[y] = '.';
         }
      /* pad last line if < 16 bytes */
         if ( y < 15)
            for (; y < 16; y++)
               sprintf(hex, "%s   ", hex);
      
         fprintf(dumpFile, "%08X : %s\t%s\n", x, hex, ascii);
      }
      fclose(dumpFile);
      return(1);
   }

   int SaveAsDiff(char *filename){
      return(0);
   }
   int SaveAsBinary(char *filename){
      return(0);
   }
