//
//            MyInfo Plugin for Ollydbg
//                        flyinghearts <flyinghearts@qq.com>
// 
//  This program is free software; you can redistribute it and/or modify it
//  under the terms of the GNU General Public License as published by the 
//  Free Software Foundation; 
//
//  label/comment file format:
//    L/C address name
//  L/C: label/Comment
//
//  break point file format (.bp file):
//   B "fn" A/R address E/D passcount comment S condition S expression S explanation S plugcmd N
//
//   B: break point
//   fn: filename
//   A/R: the follow address is Absolute/Relative address
//   E/D: Enable/Disable
//   S: '\x1e' (you could change Sep defined below)
//   N: "\x1e\n"
//   passcount/condition/expression/explanation/plugcmd: 
//       see:  Conditional logging breakpoint dialog (Shift+F4)
//     Notice: There are two characters, '\x1e' and a space/tab, between each.  
//
//  example: 
//   B "abc.exe" A 00402253 E 0
//   B "abc.dll" R 0000108A D 0
//   B "test.exe" A 00402265 E 0 This is a comment.
//   B "test.exe" A 0040227F E 0 a comment eax == 0
//

#include <windows.h>
#include <stdio.h>
#include <string.h>
#include "plugin.h"

#ifdef _MSC_VER
  #pragma comment(lib,"OLLYDBG.LIB")
  #define snprintf _snprintf
  #define strnicmp  _strnicmp
#endif

#define Cache_size 256

//Modify the following values to change default behaviour
#define BP_EXT ".bp"    //break points
#define LC_EXT ".txt"   //Label/Comment
//delimeter used in functions for importing/exporting breakpoints,
//don't change it to a space/tab.
#define Sep "\x1e" 

void copy_labels_to_comment();
void export_name(char cmd);
void import_name();
void delete_extra_comment();
void import_bp();
void export_bp();
void delete_all_bps();

HINSTANCE hinst;
HWND      hwmain;

BOOL WINAPI DllEntryPoint(HINSTANCE hi,DWORD reason,LPVOID reserved){
  if (DLL_PROCESS_ATTACH == reason) hinst = hi;
  return 1;
}

extc int _export cdecl ODBG_Plugindata(char shortname[32]){
  strcpy(shortname,"MyInfo");
  return PLUGIN_VERSION;
}

extc int _export cdecl ODBG_Plugininit(int ollydbgversion,HWND hw,ulong *features){
  if (ollydbgversion < PLUGIN_VERSION) return -1;
  hwmain = hw;
  Addtolist(0,0, "MyInfo plugin v1.2 by flyinghearts");
  return 0;
}

extc int _export cdecl ODBG_Pluginmenu(int origin,char data[4096],void *item){
  switch (origin){
    case PM_MAIN:
      strcpy(data,"1 &1 Copy user-defined labels to user-defined comments|"
                  "2 &2 Export user-defined labels to file|"
                  "3 &3 Export user-defined comments to file|"
                  "4 &4 Import labels/comments from file|"              
                  "5 &5 Delete extra comments|"              
                  "7 &7 Export breakpoints|"
                  "8 &8 Import breakpoints|"
                  "9 &9 delete all breakpoints|"
                  "0 &0 About"
            );
      return 1;
   
    case PM_BREAKPOINTS:
      if (Plugingetvalue(VAL_MAINBASE) == 0) return 0;
      strcpy(data, "#MyInfo {"
                  "1 &1 Export breakpoints|"
                  "2 &2 Import breakpoints|"
                  "3 &3 Delete all breakpoints"
                  "}"
            );
      return 1; 
      
    default: break;
  }
  return 0;
}

extc void _export cdecl ODBG_Pluginaction(int origin,int action,void *item) {
  if (PM_MAIN == origin){
    switch (action){
      case 1: copy_labels_to_comment(); break;     //copy labels to comments
      case 2: export_name('L'); break;     //export labels
      case 3: export_name('C'); break;     //export comments
      case 4: import_name(); break;        
      case 5: delete_extra_comment(); break;    
      case 7: export_bp(); break;     
      case 8: import_bp(); break;     
      case 9: delete_all_bps(); break;     
      case 0: MessageBox(hwmain, "MyInfo plugin v1.2 " 
                "\n  by flyinghearts (flyinghearts@qq.com)\n\n",
                "About",MB_OK|MB_ICONINFORMATION);
               break;  
      default: break;
    }
  } else if(PM_BREAKPOINTS == origin){
    switch (action){
      case 1: export_bp(); break;   
      case 2: import_bp(); break;
      case 3: delete_all_bps(); break;
      default: break;      
    }
  }  
}


void copy_labels_to_comment()
{
  ulong addr,start,start_code,end_code;
  ulong first = 0, last = 0, count = 0;
  char name[TEXTLEN]; 
  char result[Cache_size];
  int ret;
  BOOL check_code = 1;
  struct t_module *module;
  struct t_dump *dump;
  
  dump = (t_dump*)Plugingetvalue(VAL_CPUDASM);
  if (dump == NULL || (start = dump->base) == 0 
      || (module = Findmodule(start)) == 0) {
    MessageBox(hwmain, "Failed!","Failed",MB_OK);
    return;
  }
  start = module->base;
  start_code = module->codebase;
  end_code = start_code + module->codesize;
  Findname(start-1,NM_LABEL,name); 
  
  while ((addr = Findnextname(name)) != 0) {
    if (check_code && (addr >= end_code || addr < start_code)) {   
      ret = MessageBox(hwmain,"There are some labels outside the code"
                       " segment.\nWhether to copy them or not?\n",
                       "Choose",MB_YESNO|MB_DEFBUTTON2);    
      if (IDNO == ret) break;
      check_code = 0;
    }
    
    if (Quickinsertname(addr, NM_COMMENT, name) == -1) {
      snprintf(result,Cache_size,"Failed to insert name at address %lX!",addr);
      MessageBox(hwmain, result, "Failed", MB_OK);     
      continue;
    }
    
    if (0 == first) first = addr;
    last = addr;
    ++count;    
  }
  
  Mergequicknames();
  Broadcast(WM_USER_CHALL,0,0);
  snprintf(result,Cache_size, "%d names from 0x%lX-0x%lX are copied.\n",
           count, first, last);      
  MessageBox(hwmain,result,"Result",MB_OK);
}

void delete_extra_comment()
{
  ulong start, end, start_code, end_code;
  ulong first = 0, last = 0, count = 0;
  int ret;
  BOOL check_code = 1;
  struct t_module *module;
  struct t_dump *dump;
  
  dump = (t_dump*)Plugingetvalue(VAL_CPUDASM);
  if (dump == NULL || (start = dump->base) == 0 
      || (module = Findmodule(start)) == 0) {
    MessageBox(hwmain, "Failed!","Failed",MB_OK);
    return;
  }
  
  ret = MessageBox(hwmain, "Comments not in code segment but in current module"
             ", will be deleted!\nContinue or not?",
             "Choose", MB_YESNO||MB_DEFBUTTON2);
  if (IDNO == ret) return;
  
  start = module->base;
  end = start + module->size;
  start_code = module->codebase;
  end_code = start_code + module->codesize;
  Deletenamerange(start, start_code, NM_COMMENT);
  Deletenamerange(end_code + 1, end, NM_COMMENT);
  Broadcast(WM_USER_CHALL,0,0);
}

void export_name(char cmd)
{
  ulong addr,start;
  ulong first = 0, last = 0, count = 0;
  int nametype = NM_LABEL;
  struct t_module *module;
  struct t_dump *dump; 
  FILE *fp;
  char filename[MAX_PATH];
  char name[TEXTLEN];
  char result[Cache_size];
  char type = 'L';
  char *strtype[] = {"Label","Comment"};   

  switch(cmd){
    case 'C': nametype = NM_COMMENT; 
              strtype[0] = strtype[1];
              type = 'C';               
    case 'L':      
              break;
    default: return;
  }
  
  dump = (t_dump*)Plugingetvalue(VAL_CPUDASM);
  if (dump == NULL || (start = dump->base) == 0 || (module = Findmodule(start)) == 0){
    MessageBox(hwmain, "Failed!","Failed",MB_OK);
    return;
  }
  start = module->base;

  snprintf(result,Cache_size,"Select file for %s Export!",strtype[0]);
  filename[0] = 0;
  if (Browsefilename(result, filename,LC_EXT,0) == 0) return;
  if ((fp = fopen(filename,"w")) == NULL) {
    MessageBox(hwmain, "Failed to open file to write","Failed",MB_OK);
    return;
  }

  Findname(start-1,nametype,name);
  while ((addr = Findnextname(name)) != 0) {
    fprintf(fp,"%c %lX %s\n",type,addr,name);
    if (0 == first) first = addr;
    last = addr;
    ++count;
  }
  
  fclose(fp);
  snprintf(result,Cache_size, "%d names from 0x%lX-0x%lX are exported.\n",
           count, first, last);      
  MessageBox(hwmain, result,"Result",MB_OK);
}

void import_name() 
{
  char filename[MAX_PATH];
  char name[TEXTLEN];                  //Label/Comment name
  char result[Cache_size];
  char format4[128];                    //string format with 4 args for reading file 
  char *strtype[] = {"Label","Comment"};
  char *cmdtype[] = {"Import","Export"};  
  char im_str[4] = "LC";                 //(import lables/comments)
  char newline[2];          
  ulong addr,start,end;
  ulong first = 0, last = 0, count = 0, line = 0;
  int nametype = NM_LABEL;
  int ret;
  BOOL ignore_addr = 0;
  FILE *fp;
  struct t_module *module;
  struct t_dump *dump;  
  
  //snprintf(format4,128,"%%1[LC]%%*1[ \t]%%8lX%%*1[ \t]"
  snprintf(format4,128,"%%1[LC]%%*[ \t]%%8lX%%*[ \t]"
           "%%%u[\t\x20-\xff]%%1[\n]",TEXTLEN-1);
           
  dump = (t_dump*)Plugingetvalue(VAL_CPUDASM);
  if (dump == NULL || (start = dump->base) == 0 
       || (module = Findmodule(start)) == 0) {
    MessageBox(hwmain, "Failed!","Failed",MB_OK);
    return;
  }
  start = module->base;
  end = start + module->size;

  snprintf(result,Cache_size,"Select file for %s Import",strtype[0]);
  filename[0] = 0;
  if (Browsefilename(result,filename,LC_EXT,0) == 0) return;
  if ((fp = fopen(filename,"r")) == NULL){
    MessageBox(hwmain,"Failed to open file to read!","Failed",MB_OK);
    return;
  }
  
  while (1){
    ++line;
    fscanf(fp,"%*[ \t]");
    while (fscanf(fp,"%1[\n]",newline) == 1) ++line;  //ignore blank lines
    ret = fscanf(fp,format4,im_str,&addr,name,newline);
    if (EOF == ret) break;  
    if (ret != 4){
      snprintf(result,Cache_size,"Format error at line: %d!(%d/4 matched!)\n"
              "Ignore the line and continue?\n",line,ret);
      ret = MessageBox(hwmain, result,"Choose",MB_YESNO|MB_DEFBUTTON2);      
      if (IDNO == ret) break;
      if (fscanf(fp,"%*[^\n]%1[\n]",newline) != 1) break; 
      else continue;
    } 
    nametype = 'L' == im_str[0] ? NM_LABEL : NM_COMMENT; 
    
    if (addr >= end || addr < start) {
      if (ignore_addr) continue;
      snprintf(result,Cache_size,"Line %d:\nAddress 0x%lX is not in current"
               "module(0x%lX-0x%lX)\nClick yes/no to accept/disallow all"
               "those not in the module,\nand cancel to ignore the line.)\n\n",
                line,addr,start,end-1);
      ret = MessageBox(hwmain, result,"Choose",MB_YESNOCANCEL|MB_DEFBUTTON2);  
      if (IDNO == ret) break;
      else if (IDCANCEL == ret) continue;
      else ignore_addr = 1;
    }
   
    if (Quickinsertname(addr, nametype, name) == -1){      
      snprintf(result,Cache_size,"Failed to insert name at address %lX!",addr);
      MessageBox(hwmain, result, "Failed", MB_OK);     
      continue;
    }

    if (0 == first) first = addr;
    last = addr;
    ++count;
  }
  
  fclose(fp);                 //case: import/export
  Mergequicknames();
  Broadcast(WM_USER_CHALL,0,0);
  snprintf(result,Cache_size, "%d names from 0x%lX-0x%lX are imported.\n",
           count, first, last);      
  MessageBox(hwmain, result,"Result",MB_OK);
}

void delete_all_bps()
{
  ulong main_base;
  int n, ret;
  t_table  *table;
  t_bpoint *bp;
  
  main_base = Plugingetvalue(VAL_MAINBASE);
  if (0 == main_base) return;
  table = (t_table*)Plugingetvalue(VAL_BREAKPOINTS);
  if (table == NULL || (n = table->data.n) == 0) return;
  
  ret = MessageBox(hwmain, "All int3 breakpoints will be deleted!\n"
                   "Continue or not?","Warning!",MB_YESNO|MB_DEFBUTTON2);
  if (ret == IDNO) return;
 
  bp = (t_bpoint*)Getsortedbyselection(&(table->data), 0);
  Deletebreakpoints(bp->addr, (bp + n - 1)->addr + 1,0);
  Broadcast(WM_USER_CHALL,0,0);  
}

void export_bp()
{
  char filename[MAX_PATH];
  char name[TEXTLEN];
  char expr[TEXTLEN];
  char expl[TEXTLEN];
  char comment[TEXTLEN];
  char plugcmd[TEXTLEN];
  char result[Cache_size];
  char bp_type,addr_type;
  char *module_name;
  ulong type,addr,passcount,main_base,base;
  int i,n,count, total;
  FILE  *fp;
  t_table  *table;
  t_module *module;
  t_bpoint *bp;
  
  main_base = Plugingetvalue(VAL_MAINBASE);
  if (0 == main_base) return;
  table = (t_table*)Plugingetvalue(VAL_BREAKPOINTS);
  if (table == NULL || (n = table->data.n) ==0 ) return;

  filename[0] = 0;
  if (Browsefilename("Select file for breakpoints Export.",filename,BP_EXT,0)
       == 0) return;  
  if ((fp = fopen(filename, "w")) == NULL){
    MessageBox(hwmain, "Failed to open file to write!","Failed",MB_OK);
    return;
  }
  
  bp = (t_bpoint*)Getsortedbyselection(&(table->data), 0);
  for (i = 0, count = 0; i < n; i++){
    addr = bp->addr;
    type = bp->type;
    if ((TY_INVALID | TY_ONESHOT | TY_TEMP) & type) continue;

    //If TY_ACTIVE and TY_DISABLED are set simultaneously,TY_DISABLED is ignored
    if (type & TY_ACTIVE) bp_type = 'E';        
    else if(type & TY_DISABLED) bp_type = 'D';
    else continue;

    passcount = bp->passcount;
    bp++;
    if ((module = Findmodule(addr)) == 0) continue;
    base = module->base;
    if (base == main_base) addr_type = 'A';  //absolute address
    else addr_type = 'R', addr -= base;     //relative address
    
    Findname(addr,NM_COMMENT,comment);
    total = 0;
    total |= Findname(addr,NM_BREAK,name);
    total |= Findname(addr,NM_BREAKEXPR,expr); 
    total |= Findname(addr,NM_BREAKEXPL,expl); 
    total |= Findname(addr,NM_PLUGCMD,plugcmd);    
    
/*  cond=expr[0] & COND_MASK;
    j=0;
    if(cond & COND_NOBREAK  )  bp_cond[j++]='N';
    if(cond & COND_BRKALWAYS)  bp_cond[j++]='B';
    if(cond & COND_LOGTRUE  )  bp_cond[j++]='l';
    if(cond & COND_LOGALWAYS)  bp_cond[j++]='L';
    if(cond & COND_ARGTRUE  )  bp_cond[j++]='a';
    if(cond & COND_ARGALWAYS)  bp_cond[j++]='A';
    bp_cond[j]=0; */
    
    module_name = strrchr(module->path,'\\');
    module_name = module_name == NULL ? module->path : module_name + 1;
    ++count;
    
    if (total == 0) {                        //@@@
      fprintf(fp,"B \"%s\" %c %08lX %c %lX %s\n", 
              module_name,addr_type,addr,bp_type,passcount,comment);      
    } else {
      fprintf(fp,"B \"%s\" %c %08lX %c %lX %s" Sep " %s" Sep " %s" Sep 
              " %s" Sep " %s" Sep "\n", module_name,addr_type,addr,
              bp_type,passcount,comment,name,expr,expl,plugcmd);
    }
  }
  
  fclose(fp);
  snprintf(result,Cache_size,"%d breakpoints are exported.\n",count);
  MessageBox(hwmain, result,"Export result",MB_OK);    
}

void import_bp()
{
  char filename[MAX_PATH+1];
  char name[TEXTLEN+1];
  char expr[TEXTLEN+1];
  char expl[TEXTLEN+1];
  char comment[TEXTLEN+1];
  char plugcmd[TEXTLEN+1];
  char result[Cache_size];
  char bp_type[2];
  char addr_type[2];
  char newline[2];
  char *module_name;
  char *main_name;
  char *prev_name;
  ulong type,addr,passcount,main_base,main_end,prev_base,main_size,prev_size;
  int ret,count,i,n;
  BOOL check_main_name = 1;
  FILE  *fp;
  t_table  *table;
  t_module *module, *module_s;
  
#define Mkstr(a) # a
#define Mks(b) Mkstr(b)
#define Sp "%*[ \t]"
#define Ss "%" Mks(TEXTLEN) "[^" Sep "]" Sep
#define Sr "%" Mks(TEXTLEN) "[^" Sep "\n]" Sep
  //Use %lX to get passcount value, may get unwanted result. When passcount value is not
  //given in the input line, the 'B' character at the next line may be read.
  char format10[] = "B" Sp "\"%" Mks(MAX_PATH) "[^\"]\"" Sp "%1[AR]"
                    Sp "%8lX" Sp "%1[DE]" Sp "%lu" Sr Sr Ss Sr Sr;

  main_base = Plugingetvalue(VAL_MAINBASE);
  if (main_base == 0) return;
  module = (t_module*) Findmodule(main_base);
  main_size = module->size;
  main_end = main_base + main_size;
  main_name = strrchr(module->path,'\\');
  main_name = main_name == NULL ? module->path : main_name + 1;
  prev_name = main_name;
  prev_base = main_base;
  prev_size = main_size;
  
  table = (t_table*)Plugingetvalue(VAL_MODULES);
  module_s = (t_module*) Getsortedbyselection(&(table->data), 0);
  if (table == NULL || (n = table->data.n) == 0) return;

  filename[0] = 0;
  if (Browsefilename("Select file for breakpoints Import",filename,BP_EXT,0)
       == 0) return;    
  if ((fp = fopen(filename, "r")) == NULL){
    MessageBox(hwmain, "Failed to open file to read!","Failed",MB_OK);
    return;
  }
  
  count = 0; 
  while (1){
    //There is a space/tab before the name/expr/expl/comment/plugcmd. Ignore it.
    name[1] = expr[1] = expl[1] = comment[1] = plugcmd[1] = 0;
    passcount = 0;
    addr = 0;
    fscanf(fp,"%*[\n \t]");
    ret = fscanf(fp,format10,filename,addr_type,&addr,bp_type,&passcount,
                 comment,name,expr,expl,plugcmd);
    if (EOF == ret || ret == 0) break;
    if (ret < 4 || fscanf(fp,"%1[\n]",newline) != 1) {
      if (ret != 4) {
        snprintf(result,Cache_size,"only %d matched\naddr:%lX\n",ret,addr);
        MessageBox(hwmain,result,"Error.",MB_OK);
        break;
      }
      //when ret == 4, it failed to read passcount value, but spaces,including newline, 
      //may be ignored when tried to get passcount value.      
    }
    
    if ('R' == addr_type[0]){        //case: Relative address
      if (strnicmp(filename,prev_name,MAX_PATH) != 0){  
        module = module_s;
        for (i = 0; i < n; ++i, ++module){
          module_name = strrchr(module->path,'\\');
          module_name = module_name != NULL ? module_name + 1 : module->path; 
          if (strnicmp(filename,module_name,MAX_PATH) != 0) continue;
          prev_base = module->base;
          prev_size = module->size;
          prev_name = module_name;
          break;
        }
        
        if (i == n){     //module name is not found
          snprintf(result,Cache_size,"Failed to find module: %s!",filename);      
          MessageBox(hwmain,result,"Error.",MB_OK);
          break;
        } 
      }
      
      if (addr >= prev_size) {
        snprintf(result,Cache_size,"RVA:%8lx is not in module:%s",addr,filename);
        MessageBox(hwmain,result,"Error",MB_OK); 
        break;
      }
      addr += prev_base;
    } else {                          //case: Absolute address
      if (addr < main_base || addr >= main_end){
        snprintf(result,Cache_size,"VA:%8lx is not in main module!",addr);
        MessageBox(hwmain,result,"Error",MB_OK);
        break;        
      }
      
      if (check_main_name && strnicmp(filename,main_name,MAX_PATH) != 0) {
        snprintf(result,Cache_size,"Main module name is not matched.\n"
                 "module: %s\ninput: %s\nIgnore the main module name check?\n",
                  main_name,filename);      
        ret = MessageBox(hwmain, result,"Choose",MB_YESNO|MB_DEFBUTTON2);
        if (IDNO == ret) break;
        check_main_name = 0;    
      }
    }
    
    type = ('E' == bp_type[0]) ? (TY_ACTIVE | TY_KEEPCOND)
            : (TY_DISABLED | TY_KEEPCOND);
            
    Setbreakpointext(addr,type,0,passcount);
    Quickinsertname(addr,NM_BREAK, name + 1);
    Quickinsertname(addr,NM_BREAKEXPR, expr + 1); 
    Quickinsertname(addr,NM_BREAKEXPL, expl + 1); 
    Quickinsertname(addr,NM_PLUGCMD, plugcmd + 1);
    Quickinsertname(addr,NM_COMMENT, comment + 1);
    ++count;
  }
  
  fclose(fp);
  Mergequicknames();
  Broadcast(WM_USER_CHALL,0,0);
  snprintf(result,Cache_size, "%d breakpoints are set.\n", count);      
  MessageBox(hwmain, result,"Result",MB_OK);
}

