IDA Script: Fixing 16bit pushed data segment references

A good friend has started reversing an old 16bit Borland C++ (3.1?) program, and had lots of stack push data segment offsets that were not correctly cross referencing.

After telling him the shortcuts for manually fixing the issue (press O for the data segment, or Alt-R for any segment offset), he wrote an IDC script to do it en mass.

Thus (made up example code)

push ds;
mov ax, 0x1234;
push ax

should look like:

push ds;
mov ax, ds:dword_1234;
push ax

Here’s his script:

#include <idc.idc>

static main()
{
  auto seg, loc;
  auto movloc, movtarget;
  auto xref;
  auto dsegbase;

  dsegbase = SegByName("dseg") * 16;
  Message("dsegbase=%x\\n", dsegbase);

  Message("========================================\\n");
  seg = FirstSeg();

  while(seg != BADADDR )
  {
    Message("----------------------------------------\\n");

    loc = SegStart(seg);

    if( Byte(loc) != 0xCD || Byte(loc+1) != 0x3F)
    {
      Message("Fixing indirect push [ds:xx] refs from %s\\n", SegName(seg));

      while(loc != BADADDR && loc < SegEnd(seg))
      {
        if (GetMnem(loc) != "push" || GetOpnd(loc, 0) != "ds")
        {
          loc = NextHead(loc, BADADDR);
          continue;
        }
        loc = NextHead(loc, BADADDR);

        if (GetMnem(loc) != "mov" || GetOpType(loc, 1) != o_imm)
        {
          loc = NextHead(loc, BADADDR);
          continue;
        }
      movloc = loc;
      movtarget = GetOpnd(movloc, 0);
      loc = NextHead(loc, BADADDR);

      if (GetMnem(loc) != "push" || GetOpnd(loc, 0) != movtarget)
      {
        continue;
      }

      // At this point, we know we're pushing a [ds:x] combo.
      //Message("%x: mov %s, %s\\n", movloc, movtarget, GetOpnd(movloc, 1));

      // Abort if there already exists a Dxref
      xref = Dfirst(movloc);
      if (xref != BADADDR)
      {
        continue;
      }

      Message(" Updating %s:%04x\\n", SegName(seg), (movloc - seg) & 0xffff);
      OpOff(movloc, 1, dsegbase);
      }
    }

    seg = NextSeg(seg);
  }
}