#include <exec/memory.h>
#include <exec/execbase.h>
#include <dos/filehandler.h>
#include <devices/hardblocks.h>
#include <devices/scsidisk.h>
#include <devices/trackdisk.h>

#include <proto/exec.h>
#include <proto/dos.h>

#include <MyStartup.h>
#include <MyLib.h>

/***********************************************/

#define NoBlock 0xffffffff

#define BLOCKS_NEEDED 4

/***********************************************/

extern void DriveInitCode(void);
extern void CodeEnd(void);

/***********************************************/

struct RDArgs *RDArgs;

struct RigidDiskBlock RigidDiskBlock;
struct BadBlockBlock BadBlockBlock;
struct PartitionBlock PartitionBlock;
struct FileSysHeaderBlock FileSysHeaderBlock;
struct LoadSegBlock LoadSegBlock;

UBYTE *HardBlockMap;
ULONG HardBlocksHi;

ULONG Blocks[BLOCKS_NEEDED];

ULONG RDBBlock;

struct IOStdReq IORequest;

struct
   {
      char *Drive;
   } Arguments;

unsigned char __chip DMABuffer[512];

/***********************************************/

void CloseAll(int RC)

{
   CloseDevice(&IORequest);
   if (HardBlockMap) FreeMem(HardBlockMap,HardBlocksHi+1);
   FreeArgs(RDArgs);
   exit(RC);
}

/***********************************************/

void LoadBlock(void *Buffer, int BufferSize, ULONG BlockNumber, ULONG BlockID)

{
   int i;
   ULONG SummedLongs;
   LONG CheckSum;
   ULONG *Temp;

   if (CheckSignal(SIGBREAKF_CTRL_C))
      {
         PrintFault(ERROR_BREAK,"InstallPassword");
         CloseAll(RETURN_WARN);
      }
   IORequest.io_Command=CMD_READ;
   IORequest.io_Length=TD_SECTOR;
   IORequest.io_Data=(APTR)DMABuffer;
   IORequest.io_Offset=BlockNumber*TD_SECTOR;
   if (DoIO(&IORequest))
      {
         printf("Error reading block %ld\n",BlockNumber);
         CloseAll(RETURN_ERROR);
      }
   if (Buffer)
      {
         __builtin_memcpy(Buffer,DMABuffer,BufferSize);
         if (BlockID!=((struct RigidDiskBlock *)Buffer)->rdb_ID)
            {
               printf("Block %ld: exspected type 0x%lx got type 0x%lx\n",BlockNumber,BlockID,((struct RigidDiskBlock *)Buffer)->rdb_ID);
               CloseAll(RETURN_ERROR);
            }
         SummedLongs=((struct RigidDiskBlock *)Buffer)->rdb_SummedLongs;
         CheckSum=0;
         Temp=(ULONG *)Buffer;
         for (i=0; i<SummedLongs; i++)
            {
               CheckSum=CheckSum+(*Temp++);
            }
         if (CheckSum)
            {
               printf("Checksum of block %ld is invalid.\n",BlockNumber);
               CloseAll(RETURN_ERROR);
            }
         if (HardBlockMap)
            {
               HardBlockMap[BlockNumber]=TRUE;
            }
      }
}

/***********************************************/

void PrintLoadSegBlocks(ULONG Block)

{
   while (Block!=NoBlock)
      {
         LoadBlock(&LoadSegBlock,sizeof(struct LoadSegBlock),Block,IDNAME_LOADSEG);
         Block=LoadSegBlock.lsb_Next;
      }
}

/***********************************************/

void PrintFileSysHeaderBlock(struct FileSysHeaderBlock *FHB)

{
   ULONG Block;

   if ((Block=FHB->fhb_SegListBlocks)==NoBlock) return;
   PrintLoadSegBlocks(Block);
}

/***********************************************/

void MakeHardblockMap(void)

{
   ULONG Block;
   int Count;

   Block=FALSE;
   for (RDBBlock=0; RDBBlock<RDB_LOCATION_LIMIT && !Block; )
      {
         LoadBlock(NULL,512,RDBBlock,NULL);
         if (*(ULONG *)DMABuffer==IDNAME_RIGIDDISK)
            {
               Block=TRUE;
            }
         else
            {
               RDBBlock++;
            }
      }
   if (!Block)
      {
         printf("Sorry, I couldn't find the RigidDiskBlock on this unit!\n"
                "The password code could not be installed.\n");
         CloseAll(RETURN_FAIL);
      }
   printf("Found the RigidDiskBlock on block %ld\n",RDBBlock);
   LoadBlock(&RigidDiskBlock,sizeof(struct RigidDiskBlock),RDBBlock,IDNAME_RIGIDDISK);
   if (RigidDiskBlock.rdb_DriveInit!=NoBlock)
      {
         printf("Sorry, I couldn't install the password code.\n"
                "Some DriveInit code is already installed on this unit!\n");
         CloseAll(RETURN_ERROR);
      }
   HardBlocksHi=RigidDiskBlock.rdb_RDBBlocksHi;
   if (HardBlockMap=AllocMem(HardBlocksHi+1,MEMF_CLEAR))
      {
         for (Block=0; Block<=RDBBlock; Block++)
            {
               HardBlockMap[Block]=TRUE;
            }
      }
   else
      {
         printf("Out of memory.\n");
         CloseAll(RETURN_ERROR);
      }

   Block=RigidDiskBlock.rdb_BadBlockList;
   if (Block!=NoBlock)
      {
         do
            {
               LoadBlock(&BadBlockBlock,sizeof(struct BadBlockBlock),Block,IDNAME_BADBLOCK);
               Block=BadBlockBlock.bbb_Next;
            }
         while (Block!=NoBlock);
      }

   Block=RigidDiskBlock.rdb_PartitionList;
   if (Block!=NoBlock)
      {
         do
            {
               LoadBlock(&PartitionBlock,sizeof(struct PartitionBlock),Block,IDNAME_PARTITION);
               Block=PartitionBlock.pb_Next;
            }
         while (Block!=NoBlock);
      }

   Block=RigidDiskBlock.rdb_FileSysHeaderList;
   if (Block!=NoBlock)
      {
         do
            {
               LoadBlock(&FileSysHeaderBlock,sizeof(struct FileSysHeaderBlock),Block,IDNAME_FILESYSHEADER);
               PrintFileSysHeaderBlock(&FileSysHeaderBlock);
               Block=FileSysHeaderBlock.fhb_Next;
            }
         while (Block!=NoBlock);
      }

   Block=RigidDiskBlock.rdb_DriveInit;
   if (Block!=NoBlock)
      {
         PrintLoadSegBlocks(Block);
      }

   for (Block=0; Block<=HardBlocksHi; Block++)
      {
         if (!HardBlockMap[Block])
            {
               LoadBlock(NULL,512,Block,NULL);
               if (*(ULONG *)DMABuffer==0x46494C4C)
                  {
                     HardBlockMap[Block]=TRUE;
                  }
            }
      }

   Count=0;
   for (Block=0; Block<=HardBlocksHi && Count<BLOCKS_NEEDED; Block++)
      {
         if (!HardBlockMap[Block])
            {
               Blocks[Count]=Block;
               Count++;
            }
      }
   if (Count!=BLOCKS_NEEDED)
      {
         printf("Sorry, I couldn't find a place to store the password code!\n");
         CloseAll(RETURN_ERROR);
      }
}

/***********************************************/

void OpenDriveDevice(void)

{
   struct DosList *DosList;
   struct FileSysStartupMsg *StartupMsg;
   char *Device;
   ULONG Unit;
   int RC;

   RC=0;
   DosList=LockDosList(LDF_READ | LDF_DEVICES);
   if (DosList=FindDosEntry(DosList,Arguments.Drive,LDF_DEVICES))
      {
         StartupMsg=(struct FileSysStartupMsg *)BADDR(DosList->dol_misc.dol_handler.dol_Startup);
         IORequest.io_Message.mn_ReplyPort=&((struct Process *)SysBase->ThisTask)->pr_MsgPort;
         Device=((char *)BADDR(StartupMsg->fssm_Device))+1;
         Unit=StartupMsg->fssm_Unit;
         if (OpenDevice(Device,Unit,&IORequest,0))
            {
               printf("Unable to open %s unit %lu\n",Device,Unit);
               RC=RETURN_FAIL;
            }
         else
            {
               printf("Attempting to install password code on %s unit %ld\n",Device,Unit);
            }
      }
   else
      {
         printf("\x22%s:\x22 not found.\n",Arguments.Drive);
         RC=RETURN_FAIL;
      }
   UnLockDosList(LDF_READ | LDF_DEVICES);
   if (RC) CloseAll(RC);
}

/***********************************************/

void WriteDriveInitCode(void)

{
   int Count, i;
   ULONG Code;
   long CheckSum;

   LoadBlock(&RigidDiskBlock,sizeof(struct RigidDiskBlock),RDBBlock,IDNAME_RIGIDDISK);
   Code=(ULONG)DriveInitCode;
   for (Count=0; Count<BLOCKS_NEEDED; Count++)
      {
         LoadSegBlock.lsb_ID=IDNAME_LOADSEG;
         LoadSegBlock.lsb_ChkSum=0;
         LoadSegBlock.lsb_HostID=RigidDiskBlock.rdb_HostID;
         if (Count+1!=BLOCKS_NEEDED)
            {
               LoadSegBlock.lsb_Next=Blocks[Count+1];
               LoadSegBlock.lsb_SummedLongs=TD_SECTOR>>2;
               __builtin_memcpy(LoadSegBlock.lsb_LoadData,(char *)Code,123*sizeof(ULONG));
               Code=Code+123*sizeof(ULONG);
            }
         else
            {
               LoadSegBlock.lsb_Next=NoBlock;
               LoadSegBlock.lsb_SummedLongs=(((ULONG)CodeEnd-Code+3)>>2)+5*sizeof(ULONG);
               __builtin_memcpy(LoadSegBlock.lsb_LoadData,(char *)Code,(ULONG)CodeEnd-Code);
            }
         CheckSum=0;
         for (i=0; i<LoadSegBlock.lsb_SummedLongs; i++)
            {
               CheckSum+=((long *)&LoadSegBlock)[i];
            }
         LoadSegBlock.lsb_ChkSum=0-CheckSum;
         printf("Writing block %ld\n",Blocks[Count]);
         IORequest.io_Command=CMD_WRITE;
         IORequest.io_Length=TD_SECTOR;
         IORequest.io_Data=&LoadSegBlock;
         IORequest.io_Offset=Blocks[Count]*TD_SECTOR;
         if (DoIO(&IORequest))
            {
               printf("Error writing block %ld\n",Blocks[Count]);
               CloseAll(RETURN_ERROR);
            }
      }
   RigidDiskBlock.rdb_DriveInit=Blocks[0];
   RigidDiskBlock.rdb_ChkSum=0;
   CheckSum=0;
   for (i=0; i<RigidDiskBlock.rdb_SummedLongs; i++)
      {
         CheckSum+=((long *)&RigidDiskBlock)[i];
      }
   RigidDiskBlock.rdb_ChkSum=0-CheckSum;
   printf("Writing block %ld\n",RDBBlock);
   IORequest.io_Command=CMD_WRITE;
   IORequest.io_Length=TD_SECTOR;
   IORequest.io_Data=&RigidDiskBlock;
   IORequest.io_Offset=RDBBlock*TD_SECTOR;
   if (DoIO(&IORequest))
      {
         printf("Error writing block %ld\n",RDBBlock);
         CloseAll(RETURN_ERROR);
      }
}

/***********************************************/

void main(void)

{
   int RC;

   RC=RETURN_FAIL;
   if (CommandLineLength)
      {
         if (RDArgs=ReadArgs("DRIVE/A",(long *)&Arguments,NULL))
            {
               OpenDriveDevice();
               MakeHardblockMap();
               WriteDriveInitCode();
               printf("Password V1.0 successfully installed.\n"
                      "The current password is the empty password.\n"
                      "Use the passwd command to change it.\n");
               RC=RETURN_OK;
            }
         else
            {
               PrintFault(IoErr(),"InstallPassword");
            }
      }
   CloseAll(RC);
}