/**
 **   IntuiTrack V1.0
 **
 **   (c)1994 by Matthias Meixner
 */


#include <stdio.h>
#include <exec/exec.h>
#include <proto/exec.h>
#include <exec/execbase.h>
#include <dos.h>
#include <stdlib.h>

#include <clib/alib_protos.h>

extern struct ExecBase *SysBase;
extern struct IntuitionBase *IntuitionBase;

long WatchSig;

struct Task *WatchTask;

void *MemPool;

#define BUFS 64

char buffer[BUFS][128];
int read,write;

#define TYPE_WINDOW     0
#define TYPE_SCREEN     1
#define TYPE_SCREENLOCK 2

struct ResBlock {
   struct ResBlock *Next;
   int    Type;
   void *Res;
   void *Base;
};

#define MAX_TASKS 100

void *tasklist[MAX_TASKS];
struct ResBlock *reslist[MAX_TASKS];
int taskcnt=0;


void addWatch(void *tsk)
{
   Forbid();
   if(taskcnt<MAX_TASKS-1) {
      tasklist[taskcnt]=tsk;
      reslist[taskcnt++]=NULL;
   }
   Permit();
}
void delWatch(void *tsk)
{
   int a;

   Forbid();

   for(a=0;a<taskcnt;a++) {
      if(tasklist[a]==tsk) {
         while(a<taskcnt-1) {
            tasklist[a]=tasklist[a+1];
            reslist[a]=reslist[a+1];
            a++;
         }
         break;
      }
   }

   Permit();
}
int findWatch(void *tsk)
{
   int a;
   int found=-1;

   Forbid();
   for(a=0;a<taskcnt;a++) {
      if(tasklist[a]==tsk) {
         found=a;
         break;
      }
   }
   Permit();
   return(found);
}

char *TypeStr[]={"Window","Screen","Screenlock"};

void MyAddResource(int Type,void *adr,void *base);
void MyRemResource(void *adr);
void MyFlushResource(void);

/* Intuition-patches */

void __asm *(*OldOpenScreen)( register __a0 void *ns,
                                       register __a6 void *base);

void __interrupt __saveds __asm *NewOpenScreen(register __a0 void *ns,
                                                  register __a6 void *base)
{
   void *scr;

   scr=OldOpenScreen(ns,base);

   if(scr) MyAddResource(TYPE_SCREEN,scr,base);

   return(scr);
}
void __asm *(*OldOpenWindow)( register __a0 void *nw,
                                       register __a6 void *base);

void __interrupt __saveds __asm *NewOpenWindow(register __a0 void *nw,
                                                  register __a6 void *base)
{
   void *win;

   win=OldOpenWindow(nw,base);

   if(win) MyAddResource(TYPE_WINDOW,win,base);

   return(win);
}

void __asm *(*OldOpenScreenTagList)( register __a0 void *ns,
                                     register __a1 void *tags,
                                     register __a6 void *base);

void __interrupt __saveds __asm *NewOpenScreenTagList(register __a0 void *ns,
                                                  register __a1 void *tags,
                                                  register __a6 void *base)
{
   void *scr;

   scr=OldOpenScreenTagList(ns,tags,base);

   if(scr) MyAddResource(TYPE_SCREEN,scr,base);

   return(scr);
}
void __asm *(*OldOpenWindowTagList)( register __a0 void *nw,
                                     register __a1 void *tags,
                                     register __a6 void *base);

void __interrupt __saveds __asm *NewOpenWindowTagList(register __a0 void *nw,
                                                  register __a1 void *tags,
                                                  register __a6 void *base)
{
   void *win;

   win=OldOpenWindowTagList(nw,tags,base);

   if(win) MyAddResource(TYPE_WINDOW,win,base);

   return(win);
}

BOOL __asm (*OldCloseScreen)( register __a0 void *ns,
                             register __a6 void *base);

BOOL __interrupt __saveds __asm NewCloseScreen(register __a0 void *ns,
                                               register __a6 void *base)
{
   if(ns) MyRemResource(ns);

   return(OldCloseScreen(ns,base));
}
void __asm *(*OldCloseWindow)( register __a0 void *nw,
                               register __a6 void *base);

void __interrupt __saveds __asm NewCloseWindow(register __a0 void *nw,
                                               register __a6 void *base)
{
   if(nw) MyRemResource(nw);

   OldCloseWindow(nw,base);
}

void __asm *(*OldLockScreen)( register __a0 void *ns,
                                       register __a6 void *base);

void __interrupt __saveds __asm *NewLockScreen(register __a0 void *ns,
                                                  register __a6 void *base)
{
   void *scr;

   scr=OldLockScreen(ns,base);

   if(scr) MyAddResource(TYPE_SCREENLOCK,scr,base);

   return(scr);
}
void __asm (*OldUnlockScreen)( register __a0 void *ns,
                              register __a1 void *ns2,
                             register __a6 void *base);

void __interrupt __saveds __asm NewUnlockScreen(register __a0 void *ns,
                                               register __a1 void *ns2,
                                               register __a6 void *base)
{
   void *res;
   if(ns) {
      res=OldLockScreen(ns,base);
      OldUnlockScreen(NULL,res,base);
      if(res) MyRemResource(res);
   } else {
      if(ns2) MyRemResource(ns2);
   }
   OldUnlockScreen(ns,ns2,base);
}

/* Exec-patches */

APTR __asm (*OldAddTask)(register __a1 struct Task *,
                        register __a2 APTR ,
                        register __a3 APTR , register __a6 Base);

APTR __interrupt __saveds __asm NewAddTask(register __a1 struct Task * tsk,
                        register __a2 APTR lpc,
                        register __a3 APTR fpc, register __a6 Base)
{
   Forbid();
   sprintf(buffer[write],"Task captured : %s 0x%08lx",tsk->tc_Node.ln_Name,tsk);
   write++;
   write&=BUFS-1;
   Permit();
   addWatch(tsk);
   Signal(WatchTask,1<<WatchSig);
   return(OldAddTask(tsk,lpc,fpc,Base));
}
void __asm (*OldRemTask)(register __a1 struct Task *, register __a6 Base);

void __interrupt __saveds __asm NewRemTask(register __a1 struct Task *tsk,
                                       register __a6 Base)
{
   struct Task *rtask;
   rtask=tsk?tsk:FindTask(0);

   MyFlushResource();

   if(findWatch(rtask)>=0) {

      Forbid();
      sprintf(buffer[write],"Task released : %s 0x%08lx",rtask->tc_Node.ln_Name,rtask);

      write++;
      write&=BUFS-1;
      Permit();

      delWatch(rtask);
      Signal(WatchTask,1<<WatchSig);
   }
   OldRemTask(tsk,Base);
}

/* Management functions */

void MyAddResource(int Type,void *adr,void *base)
{
   int tnum;
   struct Task *task=FindTask(0);
   struct ResBlock *res;

   if(res=LibAllocPooled(MemPool,sizeof(struct ResBlock))) {

      res->Type=Type;
      res->Res=adr;
      res->Base=base;

      if((tnum=findWatch(task))>=0) {
         res->Next=reslist[tnum];
         reslist[tnum]=res;
      }
   }

/*   Forbid();
   sprintf(buffer[write],"%s: AddResource %s  0x%08x",task->tc_Node.ln_Name,
                         TypeStr[Type],adr);
   write++;
   write&=BUFS-1;
   Permit();
   Signal(WatchTask,1<<WatchSig); */
}
void MyRemResource(void *adr)
{
   int tnum;
   struct Task *task=FindTask(0);

   if((tnum=findWatch(task))>=0) {
      struct ResBlock *res,*mem=NULL;

      if(!(res=reslist[tnum])) return;

      if(reslist[tnum]->Res==adr) {
         mem=reslist[tnum];
         reslist[tnum]=reslist[tnum]->Next;
      } else {
         while(res->Next) {
            if(res->Next->Res==adr) {
               mem=res->Next;
               res->Next=res->Next->Next;
               break;
            }
            res=res->Next;
         }
      }

      if(mem) LibFreePooled(MemPool,mem,sizeof(struct ResBlock));
   }
}
void MyFlushResource(void)
{
   int tnum;
   struct Task *task=FindTask(0);

   if((tnum=findWatch(task))>=0) {

      struct ResBlock *res;

      while(res=reslist[tnum]) {
         reslist[tnum]=res->Next;

         switch(res->Type) {
            case TYPE_WINDOW: OldCloseWindow(res->Res,res->Base);
                              break;
            case TYPE_SCREEN: OldCloseScreen(res->Res,res->Base);
                              break;
            case TYPE_SCREENLOCK:
                              OldUnlockScreen(NULL,res->Res,res->Base);
                              break;
         }
         Forbid();
         sprintf(buffer[write],"\033[33mUnfreed %s\033[0m Task: \033[32m%s\033[0m",
                               TypeStr[res->Type],task->tc_Node.ln_Name);
         write++;
         write&=BUFS-1;
         Permit();
         Signal(WatchTask,1<<WatchSig);
         LibFreePooled(MemPool,res,sizeof(struct ResBlock));
      }

      reslist[tnum]=NULL;
   }
}

/* Safe removing of patches */

void ResetFunction(struct Library *lib,LONG offset,APTR function,APTR oldfunction)
{
   ULONG *func=(ULONG *)(((ULONG)lib)+offset+2);

   Forbid();

   while(*func!=(ULONG)function) {
      Permit();
      puts("could not remove patches ( ^C: retry )");
      Wait(SIGBREAKF_CTRL_C);
      puts("retrying ...");
      Forbid();
   }

   SetFunction(lib,offset,oldfunction);
   Permit();
}

main()
{
   long mask;

   if(!(MemPool=LibCreatePool(MEMF_PUBLIC,1024,128))) {
      fprintf(stderr,"Could not allocate memorypool\n");
      exit(EXIT_FAILURE);
   }

   WatchSig=AllocSignal(-1);
   WatchTask=FindTask(0);

   puts("IntuiTrack V1.0 ©1994 by Matthias Meixner");

   /* Exec-patches */
   OldAddTask=SetFunction((void *)SysBase,-282,(ULONG (*)()) NewAddTask);
   OldRemTask=SetFunction((void *)SysBase,-288,(ULONG (*)()) NewRemTask);

   /* Intuition-patches */
   OldOpenWindow=SetFunction((void *)IntuitionBase,-0xcc,(ULONG (*)()) NewOpenWindow);
   OldOpenScreen=SetFunction((void *)IntuitionBase,-0xc6,(ULONG (*)()) NewOpenScreen);
   OldOpenWindowTagList=SetFunction((void *)IntuitionBase,-0x25e,(ULONG (*)()) NewOpenWindowTagList);
   OldOpenScreenTagList=SetFunction((void *)IntuitionBase,-0x264,(ULONG (*)()) NewOpenScreenTagList);
   OldCloseWindow=SetFunction((void *)IntuitionBase,-0x48,(ULONG (*)()) NewCloseWindow);
   OldCloseScreen=SetFunction((void *)IntuitionBase,-0x42,(ULONG (*)()) NewCloseScreen);
   OldLockScreen=SetFunction((void *)IntuitionBase,-0x1fe,(ULONG (*)()) NewLockScreen);
   OldUnlockScreen=SetFunction((void *)IntuitionBase,-0x204,(ULONG (*)()) NewUnlockScreen);

   while(1) {
      mask=Wait((1<<WatchSig) | SIGBREAKF_CTRL_C);
      if(mask & ~SIGBREAKF_CTRL_C) {
         while(read!=write) {
            puts(buffer[read]);
            read++;
            read &=(BUFS-1);
         }
      }
      if(mask & SIGBREAKF_CTRL_C) break;
   }

   /* Exec-patches */
   ResetFunction((struct Library *)SysBase,-282,(ULONG (*)()) NewAddTask,(ULONG (*)()) OldAddTask);
   ResetFunction((struct Library *)SysBase,-288,(ULONG (*)()) NewRemTask,(ULONG (*)()) OldRemTask);

   /* Intuition-patches */
   ResetFunction((void *)IntuitionBase,-0xcc,(ULONG (*)()) NewOpenWindow,(ULONG (*)()) OldOpenWindow);
   ResetFunction((void *)IntuitionBase,-0xc6,(ULONG (*)()) NewOpenScreen,(ULONG (*)()) OldOpenScreen);
   ResetFunction((void *)IntuitionBase,-0x25e,(ULONG (*)()) NewOpenWindowTagList,(ULONG (*)()) OldOpenWindowTagList);
   ResetFunction((void *)IntuitionBase,-0x264,(ULONG (*)()) NewOpenScreenTagList,(ULONG (*)()) OldOpenScreenTagList);
   ResetFunction((void *)IntuitionBase,-0x48,(ULONG (*)()) NewCloseWindow,(ULONG (*)()) OldCloseWindow);
   ResetFunction((void *)IntuitionBase,-0x42,(ULONG (*)()) NewCloseScreen,(ULONG (*)()) OldCloseScreen);
   ResetFunction((void *)IntuitionBase,-0x1fe,(ULONG (*)()) NewLockScreen,(ULONG (*)()) OldLockScreen);
   ResetFunction((void *)IntuitionBase,-0x204,(ULONG (*)()) NewUnlockScreen,(ULONG (*)()) OldUnlockScreen);

   FreeSignal(WatchSig);
   puts("All patches removed.");
   LibDeletePool(MemPool);
}