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

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


/* 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[PATH_MAX];
   FILE *fTemp;
   FILE *tmp;
   DIR *dirBDB;
   struct stat tmpStat;

   DEBUG_PrintMsg("Loading TARGET filename: %s\n", targetName);
//    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));   
   }

//  Make sure no DB is currently loaded
   if (disasm_env.flags) CloseDB();

   /* If .bdb or .target.bdb/ exists, use those instead */
   sprintf(target.dbname, "%s.bdb", target.name);
   if (! stat(target.dbname, &tmpStat)  )   {
      /* target DB already exists: open it */
      return(LoadBDB(target.dbname));
   }
   sprintf(cmdBuf, "%s/.%s", disasm_env.dbpath, target.dbname);
   if ( ! stat(cmdBuf, &tmpStat)      )   {
      return(LoadBDB(cmdBuf));
   }
   /* If $DB_PATH/target.bdb exists, use that instead */
   sprintf(cmdBuf, "%s/%s", disasm_env.dbpath, target.dbname);
   if (! stat(cmdBuf, &tmpStat)){
      return(LoadBDB(cmdBuf));
   }

   /* from here on we are making a new DB */
   LoadDB("bdb", target.dbname, db_target);
   SwitchDB(db_target->db_id);

//   Copy target binary to .bdb
   sprintf(cmdBuf, "cp -f %s%s %s/.%s.bdb/%s", target.path, 
      target.name, disasm_env.dbpath, target.name, target.name);
   if(!((tmp=popen(cmdBuf, "r")) && (pclose(tmp)==0))) {
      return(SetLastError(4540));
   }

//   memory map target image
   sprintf(cmdBuf, "%s/.%s.bdb/%s", disasm_env.dbpath, target.name, 
            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);

   if (target.image == (void *)-1) {
      target.image = 0;
      return(SetLastError(9030));
   }
//    Create a dummy address in DB to represent entire file
   DefineAddr( 0, target.size, 0, ADDR_DATA);

   disasm_env.flags |= DB_LOADED || DB_MOD;
   target.status = DISASM_TGT_LOADED;
   initSequences();
   return(1);
}

int LoadBDB(char *dbname) {
/* Load a previously-saved .bdb 
   FILE *tmp;
 */
   struct stat s;
   FILE *f;
   FILE *tmp;
   int idx;
   char *pos, cmdBuf[256];

//  Make sure no DB is currently loaded 
   if (disasm_env.flags & DB_LOADED) CloseDB();

   LoadDB("bdb", dbname, db_target);
   SwitchDB(db_target->db_id);

//  Load file header
   pos = strrchr(dbname, '/');
   idx = (pos == 0) ? 0 :  pos - dbname + 1;
   sprintf(cmdBuf, "%s/.%s/header.txt", disasm_env.dbpath, &dbname[idx]);
   if ( ! stat(cmdBuf, &s)) {
      f = fopen( cmdBuf, "r");
      if (f) {
         target.header = calloc(s.st_size, 1);
         fread( target.header, s.st_size, 1, f);
         fclose(f);
      }
   }
// Fill Target Struct
   sprintf(cmdBuf, "%s/.%s/.info", disasm_env.dbpath, &dbname[idx]);
   f = fopen( cmdBuf, "r");
   if (f) {
      fscanf( f, "%s\n%s\n%d\n%ld\n%s\n%s\n%d\n%d\n", target.name, target.path,
               &target.size, &target.entry, target.arch, target.format,
               &target.status, &target.endian );
      fclose(f);
   } else return(SetLastError(4561));


   sprintf(cmdBuf, "%s/.%s/%s", disasm_env.dbpath, &dbname[idx], target.name);
   if ( ! stat(cmdBuf, &s) && target.size == s.st_size) {
      ; /* this is normal */
   } else {
      ErrorMsg(9030, "Bad .bdb file! \n");
      CloseDB();
      return(SetLastError(9030));
   } 
         
         
   target.FD = open(cmdBuf, O_RDWR);
   if (target.FD == -1) { 
      sprintf(cmdBuf, "Cannot open file %s", errno);
      ErrorMsg(9030, cmdBuf);
      return(SetLastError(9030));
   }

   target.image = mmap(0, target.size, PROT_READ | PROT_WRITE, 
      MAP_SHARED, target.FD, 0);
   if (target.image == (void *)-1) {
      target.image = 0;
      return(SetLastError(9030));
   }
   disasm_env.flags |= TGT_MAPPED;

   if (! target.header ) {
      target.header = calloc( 64, 1);
      sprintf(target.header, "Header not available!\n");
   }
   strncpy(target.dbname, dbname, PATH_MAX);

   initSequences();
   disasm_env.flags |= DB_LOADED;
   return(1);
}

/* ---------------------------------------------------------------Saving */
void SaveDBAs( char *name) {
/* store database and image in .bdb file */
   char cmdBuf[256], pwd[PATH_MAX];
   FILE *f, *tmp;

   sprintf(cmdBuf, "Saving DB to %s...\n", name);
   Message(cmdBuf);
/* create .info file */
   sprintf(cmdBuf, "%s/.%s.bdb/.info", disasm_env.dbpath, target.name);
   DEBUG_PrintVar(cmdBuf, "%s");
   f = fopen( cmdBuf, "w");
   if (f) {
      fprintf( f, "%s\n%s\n%d\n%ld\n%s\n%s\n%d\n%d\n", target.name, target.path,
               target.size, target.entry, target.arch, target.format,
               target.status, target.endian );
      fclose(f);
   }
/* create header str file */
   sprintf(cmdBuf, "%s/.%s.bdb/header.txt", disasm_env.dbpath, target.name);
   DEBUG_PrintVar(cmdBuf, "%s");
   f = fopen( cmdBuf, "w");
   if (f) {
      fprintf( f, "%s", target.header);
      fclose(f);
   }
   msync(0, target.size, MS_SYNC);
   getcwd(pwd, PATH_MAX);
   chdir(disasm_env.dbpath);
      
   sprintf(cmdBuf, "tar --remove-files -zcf %s %s/.%s.bdb", name, 
                        disasm_env.dbpath, target.name);
   DEBUG_PrintVar(cmdBuf, "%s");
   if(!((tmp=popen(cmdBuf, "r")) && (pclose(tmp)==0))) {
      unlink(name); /* does this remove all contents of subdir? */
      chdir(pwd);
      ErrorMsg(4565, "Unable to save target .bdb!");
      SetLastError(4565); 
      return;
   }   
   else Message("DB saved\n");   
   chdir(pwd);
}

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

void SaveDB( ) {
   char *str, info_file[256];
   struct stat statbuf;
   FILE *f;
   int c;

   /* the first time, create .info file */
   snprintf(info_file, 256, "%s/.%s/.info", disasm_env.dbpath, target.name); 
   if ( stat( info_file, &statbuf)) {
      f = fopen( info_file, "w");
      if (f) {
         fprintf( f, "%s\n%s\n%d\n%ld\n%s\n%s\n", target.name, target.path,
               target.size, target.entry, target.arch, target.format );
         fclose(f);
      }
   }

   if ( disasm_prefs.options & 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':
            printf("Save as: ");
            fgets(str, 256, stdin);
            if ( str[strlen(str)] == '\n' ) str[strlen(str)] = '\0';
            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];
   FILE *tmp;


   DEBUG_Print("CloseDB: Saving DB\n");
   if (disasm_env.flags & DB_MOD) {         
      SaveDB();
      disasm_env.flags ^= DB_MOD;
   }
   DEBUG_Print("CloseDB: Unloading DB\n");
   if (disasm_env.flags & DB_LOADED) {
      UnloadDB(db_target);
      disasm_env.flags ^= DB_LOADED;
   }
   DEBUG_Print("CloseDB: UnMapping Target File\n");
   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);
   DEBUG_Print("CloseDB: Switch DB\n");
   SwitchDB(db_config->db_id);
   DEBUG_Print("CloseDB: rm -rf'ing target dir\n");
   sprintf(cmdBuf, "rm -rf %s/.%s.bdb", disasm_env.dbpath, target.name);
   DEBUG_PrintVar(cmdBuf, "%s");
   if(!((tmp=popen(cmdBuf, "r")) && (pclose(tmp)==0))) {
      ErrorMsg(4511, "Error cleaning up .bdb");
   }
   if (! disasm_env.flags & QUIT_NOW ) {
      DEBUG_Print("CloseDB:  ReInit\n");
      ReInit(disasm_env.home);
   }
   DEBUG_Print("CloseDB: done\n");
   return;
}

int SaveAsASM(char *filename){
   if (! filename ) 
      return(SetLastError(9010));
   return(ext_gen_asm_file(filename));
}

int SaveAsHLL(char *filename){
   if (! filename ) 
      return(SetLastError(9010));
   return(ext_gen_hll_file(filename));
}

int SaveAsLST(char *filename){
   FILE *f;
   char *line;
   int cont;
   struct address addr;
   struct section sec;
   void *state;

   if (! filename ) 
      return(SetLastError(9010));
   f = fopen(filename, "w");
   if ( ! f) 
      return(SetLastError(9030));
   if (! (line = calloc(2048,1)) ) return(SetLastError(9050));

   fprintf( f, "Disassembly of %s by the bastard\n\n", target.name);
   fprintf( f, "%s", target.header);

   cont = DBIndexFirst(SECTION_RVA, &sec);
   while( cont ) {
      if ( (cont = DBFindClosestNext(ADDRESS_RVA, &sec.rva, &addr)) ) {
         fprintf(f, "Section: %s\tStart Address: %08X\tEnd Address: %08X\n",
                    sec.name, sec.rva, sec.rva + sec.size);
         fprintf(f, "---------------------------------------------------------"
                    "----------------\n");
         while( cont && addr.rva < sec.rva + sec.size){
            state = DBSaveState();
            ext_sprint_code(addr.rva, line, 2048, disasm_env.output);
            fprintf(f, "%s\n", line);
            line[0] = '\0';
            DBRestoreState(state);
            cont = DBIndexNext(ADDRESS_RVA, &addr);
         }
         fprintf(f, "\n\n");
      }
      cont = DBIndexNext(SECTION_RVA, &sec);
   }
   free(line);
   fclose(f);
   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);
}

int SetTargetFormat(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 libname[PATH_MAX];
   struct stat statbuf;

   DEBUG_PrintMsg("Setting TARGET format:  %s\n", file_format);
   /* if .bc not exist, assume .so */
   sprintf(libname, "%s/formats/%s.bc", disasm_env.home, file_format);
   if ( stat(libname, &statbuf) )
      /* assume shared library ... LoadExtension will catch missing .so */
      sprintf(libname, "%s/formats/lib%s.so", disasm_env.home, file_format);
   if ( ext_format->ext.filename && 
        strcmp(libname, ext_format->ext.filename) ) { 
      //UnloadExtension( EXT_FORMAT );
      //return(1);
      ;
   }
DEBUG_PrintVar(libname, "%s");
   if ( (ext_format->ext.flags & EXT_STATUS_MASK) != EXT_INIT )
       if (! LoadExtension(EXT_FORMAT, libname, ext_format) )
          return(SetLastError(4594));

   strncpy(target.format, file_format, 32);
   target.status = DISASM_TGT_FORMAT_SET;
   return(1);
}

int ApplyTargetFormat(){
   /* call the read_header routine */
   ext_read_header();
   disasm_env.flags |= DB_MOD;
   /* move this into a master PhasesOfDisassembly controller */
   target.status = DISASM_TGT_PREDISASM;
   return(1);

}

int SetTargetArch(char *file_arch) {
   char libname[PATH_MAX];
   struct stat tmpStat;

   DEBUG_PrintMsg("Setting TARGET arch:  %s\n", file_arch);
   /* default to the current disassembler */
   if ( file_arch == NULL) strcpy(libname, ext_arch->ext.filename);
   else {
      /* if the current disassembler != disasm, unload it */
      sprintf(libname, "%s/arch/%s.bc", disasm_env.home, file_arch);
      if ( stat(libname, &tmpStat) )
         /* if .bc not exist, assume .so */
         sprintf(libname, "%s/arch/lib%s.so", disasm_env.home, file_arch);
   }

   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 the disassembler isn't loaded, load it */
   if ( ext_arch->ext.flags & EXT_STATUS_MASK != EXT_INIT )
      if(! LoadExtension(EXT_ARCH, libname, ext_arch) )
         return(SetLastError(4291));

   target.endian = ext_arch->endian;
   disasm_env.flags |= DB_MOD;
   return(1);
}
int SetTargetAssembler( char *asm_output ){
   char libname[PATH_MAX];
   struct stat statbuf;

   DEBUG_PrintMsg("Setting TARGET assembler: %s\n", asm_output);
   /* if .bc not exist, assume .so */
   sprintf(libname, "%s/asm/%s.bc", disasm_env.home, asm_output);
   if ( stat(libname, &statbuf) )
      /* assume shared library ... LoadExtension will catch missing .so */
      sprintf(libname, "%s/asm/lib%s.so", disasm_env.home, asm_output);

   if ( ext_asm->ext.filename && strcmp(libname, ext_asm->ext.filename) ) {
      ; //Already loaded: do nothing
   } else if ( ext_asm->ext.flags & EXT_STATUS_MASK != EXT_INIT )
      if (! LoadExtension(EXT_ASM, libname, ext_asm) )
         return(SetLastError(4592));

   /* this mechanism will be replaced by something better.... */
   if (GetOption(COLOR_TTY)){
      SetTtyAsmFmt(ext_asm->asm_ttyColor);
      SetTtyDataFmt(ext_asm->data_ttyColor);
   } 
   target.status = DISASM_TGT_ASM_SET;
   return(1);
}

int SetTargetLanguage( char *language ){
   char libname[PATH_MAX];
   struct stat statbuf;

   DEBUG_PrintMsg("Setting TARGET Lang: %s\n", language);   
   /* if .bc not exist, assume .so */
   sprintf(libname, "%s/lang/%s.bc", disasm_env.home, language);
   if ( stat(libname, &statbuf) )
      /* assume shared library */
      sprintf(libname, "%s/lang/lib%s.so", disasm_env.home, language);

   if ( ext_lang->ext.filename && strcmp(libname, ext_lang->ext.filename) )  
      //UnloadExtension( EXT_LANG );
      return(1);
   if ( ext_lang->ext.flags & EXT_STATUS_MASK != EXT_INIT )
       if (! LoadExtension(EXT_LANG, libname, ext_lang) ) 
            return(SetLastError(4595));

   /* TODO: need some method of determining target bitness */
   ext_add_data_types( 4 );
   target.status = DISASM_TGT_LANG_SET;
   return(1);
}

int LoadPlugin( char *name ){
   char libname[PATH_MAX];
   struct stat statbuf;

   /* if .bc not exist, assume .so */
   sprintf(libname, "%s/plugins/%s.bc", disasm_env.home, name);
   if ( stat(libname, &statbuf) )
      /* assume shared library */
      sprintf(libname, "%s/plugins/lib%s.so", disasm_env.home, name);

   if (ext_plugin->ext.filename && strcmp(libname, ext_plugin->ext.filename))
      return(1);
   if (! LoadExtension(EXT_PLUGIN, libname, ext_plugin) ) 
      return(SetLastError(4596));

   return(1);
}
int RunPlugin(void *param ){
   return( ext_plugin_main( param ) );
} 

int UnloadPlugin( ){
   return(UnloadExtension(EXT_PLUGIN));
}

