#include <stdio.h>
#include <stdlib.h>     /* pour realloc() */
#include <string.h>     /* pour strdup() */
#include <dos.h>
#include <sys/stat.h>
#include "type.h"

/* Teste la presence des divers disques ou disquettes du systeme, 
        grace aux fonctions du BIOS. Liste les fichiers a la racine de
        chaque disque, s'interresse a toute les partitions. Finallement
        ecrit le fichier specifie dans le secteur de boot en le tronquant
        a 512 octets .

        Auteur: Benoit Papillault
        Creation: Samedi 5 Juillet 1997
        Derniere modification: Samedi 31 Octobre 1998

        Historique:
        10/07/1997: elimination des \r\n a la suite de la chaine "oui"

        31/10/1998: elimination de la fonction get_disk_geometry() qui ne
        fonctionnait pas correctement sous Windows 95. Amelioration de l'affichage.
        (tabulation, centrage, alignements). Aucun affichage de debugage n'est realise.

        English translation:
        test drive presence or floppy in the system by using BIOS fonctions.
        list files at the root directory of each drive and each partitions.
        At last, write the specified file into the boot sector by truncating
        it to 512 bytes.

        Most english explanation can be found in "util\wrb.c" for writing
        the boot sector file onto the correct drive and in "fat\fat2.c"
        for FAT12/16 filesystems details.
*/

#define SECTOR_SIZE     512
#define RETRY_MAX       3

#define CYL(cylsect)   ( (((cylsect) >> 8) & 0xff) | (((cylsect) & 0xc0) << 2) )
#define SECT(cylsect)    ((cylsect) & 0x3f)

#define FAT12   1       /* ces deux valeurs sont arbitraires */
#define FAT16   2

struct position
{
        word   C;      /* numero du cylindre */
        word   H;      /* numero de la tete */
        word   S;      /* numero du secteur */
};

struct partition
{
        byte    unite;          /* a utiliser avec les fonctions du BIOS */
        struct position premier_secteur;        /* le premier secteur de la partition */
        struct position dernier_secteur;        /* le dernier secteur de la partition */
        struct position disque; /* le nombre de secteurs/pistes pour 
                le disque physique */
        byte    ostype;
        char    *osname;        /* remplit lors de l'appel a analyser_partition */
        char    volume[11];        /* le nom donne avec la commande LABEL, 
                n'est valable que pour les partitions DOS. la chaine n'est
                pas terminee par 0 */
};           

/* la structure inscrite a la fin du premier secteur : 16 octets */

struct ipart
{
        byte    bootid;
        byte    beghead;
        word    begcylsect;
        byte    systid;
        byte    endhead;
        word    endcylsect;
        dword   relsect;
        dword   numsect;
};

/* contenu du premier secteur d'un disque dur, ie CHS=0/0/1 */

struct mbr
{
        byte    bootinst[446];          /* Master Block Code */
        struct ipart    parts[4];       /* la table des partitions */
        word    signature;              /* doit valoir 0xAA55 */
};

struct boot_sector
{
        byte    jmp[3];
        char    oemid[8];
        word    sector_size;    /* nombre d'octets par secteur = 512 */
        byte    cluster_size;   /* nombre de secteurs/clusteur */
        word    reserve;        /* nombre de secteurs reserves */
/* offset 0x10 */
        byte    nb_fat;         /* nombre de FAT */
        word    max_files;      /* nombre de fichiers max. a la racine */
        word    total_sector;   /* nombre de secteurs total */
        byte    media;
        word    nb_sector_fat;  /* nombre de secteurs/FAT */
        word    S;              /* nombre de secteurs/piste */
        word    H;              /* nombre de tetes */
        dword   hidden_sectors;  /* nombre de secteurs caches */
/* offset 0x20 */
        dword   total_sector2;  /* dans le cas ou total_sector = 0 */
};

struct rep_entry
{
        char    name[8];
        char    ext[3];
        byte    attr;
        byte    reserved[10];
        word    heure;
        word    date;
        word    first_cluster;
        dword   size;
};

/* constantes pour le champ attr */
#define ATTR_NORMAL     0x00
#define ATTR_RDONLY     0x01
#define ATTR_HIDDEN     0x02
#define ATTR_SYSTEM     0x04
#define ATTR_VOLUME     0x08
#define ATTR_DIRECT     0x10
#define ATTR_ARCHIV     0x20

/* quelques variables globales */
struct partition *liste_partitions = NULL;
int     nb_partitions = 0;

void afficher_ipart(struct ipart *p)
{
        char chs1[50], chs2[50];

        sprintf(chs1,"%d/%d/%d",CYL(p->begcylsect),p->beghead,
                SECT(p->begcylsect));
        sprintf(chs2,"%d/%d/%d",CYL(p->endcylsect),p->endhead,
                SECT(p->endcylsect));

        printf("%02x %10s %02x %10s %8ld %8ld\n",p->bootid,
                chs1,p->systid,chs2,p->relsect,p->numsect);
}

void afficher_position(struct position *p)
{
        char chs[50];

        sprintf(chs,"%d/%d/%d",p->C, p->H, p->S);
        printf("%10s",chs);
}

void ajouter_partition(struct partition *p)
{
        int i;

        i = nb_partitions;
        nb_partitions++;
        liste_partitions = (struct partition *)realloc(liste_partitions,
                nb_partitions*sizeof(struct partition));
        liste_partitions[i] = *p;
        liste_partitions[i].osname = strdup(p->osname);
}

void liberer_partitions(void)
{
        int i;

        for (i=0;i<nb_partitions;i++)
                free(liste_partitions[i].osname);
        free(liste_partitions);

        liste_partitions = NULL;
        nb_partitions = 0;
}

/* Note: phys2lo() et lo2phys() n'utilisent pas p->disque.C, on doit donc pouvoir
        se passer de la fonction get_disk_geometry() qui ne fonctionnent pas
        correctement sous Windows 95
*/

/* retourne le numero logique du secteur depuis le debut de la partition
        ou -1 en cas d'erreur */

dword phys2lo(struct partition *p,struct position *pos)
{
        dword first,result,last;

        first = (p->premier_secteur.S-1) 
                + ((dword)p->premier_secteur.H * p->disque.S)
                + ((dword)p->premier_secteur.C * p->disque.S * p->disque.H);
        last = (p->dernier_secteur.S-1) + ((dword)p->dernier_secteur.H * p->disque.S)
                + ((dword)p->dernier_secteur.C * p->disque.S * p->disque.H);
        result = (pos->S-1) + ((dword)pos->H * p->disque.S) 
                + ((dword)pos->C * p->disque.S * p->disque.H);

        if (first <= result && result <= last)
                return (result-first);
        return -1;
}

/* fait l'inverse, renvoie 1 en cas de succes, 0 en cas d'echec */

int lo2phys(struct partition *p,dword lo,struct position *pos)
{
        lo += (p->premier_secteur.S-1) 
                + ((dword)p->premier_secteur.H * p->disque.S)
                + ((dword)p->premier_secteur.C * p->disque.S * p->disque.H);
        pos->S = (lo % p->disque.S) + 1;
        lo = lo / p->disque.S; /* le nombre de pistes restantes */
        pos->H = lo % p->disque.H;
        pos->C = lo / p->disque.H;

/*        if (pos->C < p->disque.C)
                return 1;
        return 0;
*/
        return 1;
}

/* renvoie la partition selectionner par l'utilisateur ou NULL */

struct partition *selectionner_partitions(void)
{
        int i;
        struct partition *p;
        char buffer[10];
        dword taille;

        printf("selectionner une partition ou 0 pour annuler:\n");
        for (i=0;i<nb_partitions;i++)
        {
                p = &liste_partitions[i];
                taille = phys2lo(p,&p->dernier_secteur) -
                                phys2lo(p,&p->premier_secteur) +1;
                if (p->unite < 0x80)
                        printf("%2d) %.11s: disquette  %d [%8s] %8.2f Mo CHS=%d/%d/%d\n",
                                i+1,p->volume,p->unite+1,p->osname,
                                (float)taille/2048.0,p->premier_secteur.C,
                                p->premier_secteur.H,p->premier_secteur.S);
                else
                        printf("%2d) %.11s: disque dur %d [%8s] %8.2f Mo CHS=%d/%d/%d\n",
                                i+1,p->volume,(p->unite-0x80)+1,p->osname,
                                (float)taille/2048.0,p->premier_secteur.C,
                                p->premier_secteur.H,p->premier_secteur.S);
        }
        printf("=> ");
        fgets(buffer,sizeof(buffer),stdin);
        i = atoi(buffer);
        printf("i=%d\n",i);
        if (i>= 1 && i<=nb_partitions)
                return &liste_partitions[i-1];
        return NULL;
}

dword get_total_sector(struct boot_sector *boot)
{
        if (boot->total_sector == 0)
                return boot->total_sector2;
        return boot->total_sector;
}

/* renvoie 1 si tous est ok, 0 sinon */

/* int get_disk_geometry(byte unite,struct position *g)
{
        word    r_ax,r_cx,r_dx;

        asm {
        push    ax
        push    bx
        push    cx
        push    dx
        push    di
        push    es

        mov     ah,08h
        mov     dl,unite
        int     13h
        mov     r_ax,ax
        mov     r_cx,cx
        mov     r_dx,dx

        pop     es
        pop     di
        pop     dx
        pop     cx
        pop     bx
        pop     ax
        }

        if (r_ax != 0)
                return 0;

        g->C = CYL(r_cx);
        g->C += 1;
        g->H = r_dx >> 8;
        g->H += 1;
        g->S = SECT(r_cx);
        return 1;
}
*/

/* lit un secteur d'un disque. secteur est un numero de secteur "logique"
        le secteur "logique" 0 est le premier secteur (secteur1/piste 0/tete 0)
        de la disquette. La formule de conversion d'un secteur physique 
        en secteur logique est la suivante :
        
        secteur lo = (S-1) + (nb secteurs/piste) * H 
                + (nb secteurs/piste) * (nb cylindres [ou faces]) * C

        avec C: Cylindre, H: Tete, S = secteur physique
                
*/

int read_sector(byte lecteur,byte *buffer,struct position *p)
{
        word le_seg,l_off;
        byte C_prime,H_prime,S_prime;
        int i;

        /* C est sur 10 bits, S sur 6 bits et H sur 8 bits, avec 
                en fait les deux bits de poids forts de C dans les
                2 bits de poids fort de S */

        if (p->C >= (1<<10))
        {
                printf("numero de cylindre trop grand (%d)\n",p->C);
                return 0;
        }

        if (p->H >= (1<<8))
        {
                printf("numero de tete trop grand (%d)\n",p->H);
                return 0;
        }

        if (p->S >= (1<<6))
        {
                printf("numero de secteur trop grand (%d)\n",p->S);
                return 0;
        }

        C_prime = p->C & 0xff;
        H_prime = p->H & 0xff;
        S_prime = ((p->C >> 2) & 0xc0) | (p->S & 0x3f);

        le_seg = FP_SEG(buffer);
        l_off = FP_OFF(buffer);

        for (i=0;i<RETRY_MAX;i++)
        {
                asm   {
                push    es
                push    ax
                push    bx
                push    cx
                push    dx

                mov     ah,2    /* fonction 2 */
                mov     al,1
                mov     ch,C_prime
                mov     cl,S_prime
                mov     dh,H_prime
                mov     dl,lecteur
                mov     bx,le_seg
                mov     es,bx
                mov     bx,l_off
                int     13h

                pop     dx
                pop     cx
                pop     bx
                pop     ax
                pop     es

                jc      bad
                }
                return 1;
bad:
        }
        return 0;
} 

/* ecrit le secteur donne par la position p, buffer est une zone de memoire
        de la taille d'un secteur (512 octets ?) */

int write_sector(byte lecteur,byte *buffer,struct position *p)
{
        word le_seg,l_off;
        byte C_prime,H_prime,S_prime;
        int i;

        /* C est sur 10 bits, S sur 6 bits et H sur 8 bits, avec 
                en fait les deux bits de poids forts de C dans les
                2 bits de poids fort de S */

        if (p->C >= (1<<10))
        {
                printf("numero de cylindre trop grand (%d)\n",p->C);
                return 0;
        }

        if (p->H >= (1<<8))
        {
                printf("numero de tete trop grand (%d)\n",p->H);
                return 0;
        }

        if (p->S >= (1<<6))
        {
                printf("numero de secteur trop grand (%d)\n",p->S);
                return 0;
        }

        C_prime = p->C & 0xff;
        H_prime = p->H & 0xff;
        S_prime = ((p->C >> 2) & 0xc0) | (p->S & 0x3f);

        le_seg = FP_SEG(buffer);
        l_off = FP_OFF(buffer);

        for (i=0;i<RETRY_MAX;i++)
        {
                asm   {
                push    es
                push    ax
                push    bx
                push    cx
                push    dx

                mov     ah,3    /* fonction 3 */
                mov     al,1
                mov     ch,C_prime
                mov     cl,S_prime
                mov     dh,H_prime
                mov     dl,lecteur
                mov     bx,le_seg
                mov     es,bx
                mov     bx,l_off
                int     13h

                pop     dx
                pop     cx
                pop     bx
                pop     ax
                pop     es

                jc      bad
                }
                return 1;
bad:
        }
        return 0;
} 

char donne_lettre(byte lecteur)
{
        if (lecteur == 0 || lecteur == 1)
                return ('A'+lecteur);
        if (lecteur >= 0x80)
                return ('C'+lecteur-0x80);
        return '?';
}

/* affiche les fichiers a la racine du disque. remplit le champ volume
        de la structure partition */

void lister_racine(byte unite,struct partition *p,
        struct boot_sector *boot,int type_fat)
{
        struct rep_entry rep[16];   /* l'equivalent d'un secteur */
        dword secteur_lo;       /* le numero du secteur logique du repertoire */
        int i,j,nb_secteurs;
        struct position pos;

        /* Ici, le premier secteur de la partition a pour numero logique 0 */

        /* printf("clusters/secteur = %d\n",boot->cluster_size); */
        secteur_lo = boot->reserve + boot->nb_fat * boot->nb_sector_fat;
        nb_secteurs = (boot->max_files * 32 + boot->sector_size -1) / boot->sector_size;
        for (i=0;i<nb_secteurs;i++)
        {
                if (!lo2phys(p,secteur_lo+i,&pos))
                {
                        printf("erreur de conversion en secteur physique\n");
                        continue;
                }
                if (!read_sector(unite,(byte *)rep,&pos))
                {
                        printf("impossible de lire le secteur logique %ld\n",secteur_lo+i);
                        continue;
                }
                for (j=0;j<16;j++)
                {       
                        /* on affiche les fichiers en eliminant les fichiers qui
                                n'existent pas, les fichiers effaces et les noms
                                de volume. on garde les repertoires */
                        if (rep[j].name[0] != 0 && (byte)rep[j].name[0] != 0xe5)
                        {
                                /* la formule suivante est due a une correction,
                                        puisque ma premiere partition contient
                                        des noms de volume avec les attributs
                                        0xf et ne sont pas valides (moins de 
                                        11 caracteres) */
                                if ((rep[j].attr & 0xf) == ATTR_VOLUME)
                                {
                                     /*   printf("volume %.11s (attr=0x%x)\n",
                                                rep[j].name,rep[j].attr);
                                     */
                                        memcpy(p->volume,rep[j].name,11);
                                }
                                else
                                {        
                                        /* printf("%.8s.%.3s (attr=0x%x)\n",
                                                rep[j].name,rep[j].ext,rep[j].attr); */
                                }
                        }
                }
        }
}

void analyser_partition(byte unite,struct partition *p,
        struct boot_sector *boot)
{
        /* printf("premier secteur CHS=%d/%d/%d\n",p->premier_secteur.C,
                p->premier_secteur.H,p->premier_secteur.S);
        printf("dernier secteur CHS=%d/%d/%d\n",p->dernier_secteur.C,
                p->dernier_secteur.H,p->dernier_secteur.S);
        printf("geometrie du disque CHS=%d/%d/%d\n",p->disque.C,
                p->disque.H,p->disque.S);
        printf("boot_sector: CHS = %d/%d/%d\n",(word)(get_total_sector(boot) / boot->H / boot->S),
                boot->H,boot->S);
        printf("ostype = 0x%x\n",p->ostype);
        */

        char buffer[20];

        p->osname = NULL;
        memcpy(p->volume,"???????????",11);
        if (p->ostype == 1 || p->ostype == 4 || p->ostype == 6)
        {
                /* printf("Partition DOS '%.8s' (hidden_sectors=%ld)\n",
                        boot->oemid,boot->hidden_sectors); */
                p->disque.H = boot->H;
                p->disque.S = boot->S;
                p->osname = boot->oemid;
                if (p->ostype == 1)
                        lister_racine(unite,p,boot,FAT12);
                else
                        lister_racine(unite,p,boot,FAT16);
        }
        else
        {
                sprintf(buffer,"ID 0x%x",p->ostype);
                p->osname = strdup(buffer);
        }
        if (p->osname == NULL)
                p->osname = "";
        ajouter_partition(p);
}

/* first_sector pointe vers une zone de SECTOR_SIZE octets */

void analyser_disque(byte unite,struct boot_sector *boot)
{
        struct partition part;
        struct position tmp;
        struct boot_sector *boot2;
        int i;

     /*   printf("analyse du disque 0x%x\n",unite); */
        part.unite = unite;
/*        if (!get_disk_geometry(unite,&geo))
        {
                printf("BIOS Int 13h/08h failed!\n");
                return;
        }
*/

/*        printf("Unite 0x%x:\n",unite);
        afficher_position(&geo);
        printf("\n");
*/
        if (unite < 0x80) /* il s'agit d'une disquette */
        {
                part.premier_secteur.C = 0;
                part.premier_secteur.H = 0;
                part.premier_secteur.S = 1;
                part.disque.C = part.dernier_secteur.C 
                        = (dword)(get_total_sector(boot) / boot->H) / boot->S;
                part.disque.H = part.dernier_secteur.H = boot->H;
                part.disque.S = part.dernier_secteur.S = boot->S;
                part.ostype = 1;

/*                printf("Disquette:\t");
                afficher_position(&part.disque);
                printf("\n");
*/
                analyser_partition(unite,&part,boot);
        }
        else
        {
                struct mbr *master = (struct mbr *)boot;
                byte buffer[SECTOR_SIZE];
                
                boot2 = (struct boot_sector *)buffer;

                for (i=0;i<4;i++)
                {
                        /* on examine master->parts[i] */
                        part.premier_secteur.C = CYL(master->parts[i].begcylsect);
                        part.premier_secteur.H = master->parts[i].beghead;
                        part.premier_secteur.S = SECT(master->parts[i].begcylsect);
                        part.dernier_secteur.C = CYL(master->parts[i].endcylsect);
                        part.dernier_secteur.H = master->parts[i].endhead;
                        part.dernier_secteur.S = SECT(master->parts[i].endcylsect);
                        part.ostype = master->parts[i].systid;

                        /* afficher_ipart(&master->parts[i]); */

                        /* on lit le premier secteur de la partition */
                        if (!read_sector(unite,buffer,&part.premier_secteur))
                        {
                                printf("erreur de lecture du premier secteur de la partition\n");
                                continue;
                        }

/*                        tmp.C = (dword)(get_total_sector(boot2) / boot2->H) / boot2->S;
                        tmp.H = boot2->H;
                        tmp.S = boot2->S;

                        part.disque.H = boot2->H;
                        part.disque.S = boot2->S;

                        printf("Partition:\t");
                        afficher_position(&tmp);
                        printf("\n");
*/
/*                        printf("Partition %d:\t",i+1);
                        printf("relsect = %ld,numsect = %ld\n",
                                master->parts[i].relsect,
                                master->parts[i].numsect);
*/
                        analyser_partition(unite,&part,boot2);
                }
        }
}

/* charge le nouveau code du boot secteur a partir du fichier specifie,
        retourne 1 si tous est ok, 0 sinon */

int charger_boot(const char *file,byte *nouveau_boot)
{
        struct stat buf;
        char buffer[256];
        FILE *fp;
        char *p;

        if (stat(file,&buf)!=0)
        {
                perror(file);
                return 0;
        }

        if (buf.st_size > SECTOR_SIZE)
        {
                printf("le fichier fait %ld octets, il sera tronque a %d octets\n",
                        buf.st_size,SECTOR_SIZE);
                printf("voulez-vous continuer (oui) ?");
                fgets(buffer,sizeof(buffer),stdin);
                p = strchr(buffer,'\r');
                if (p!=NULL) *p=0;
                p = strchr(buffer,'\n');
                if (p!=NULL) *p=0;
                if (strcmp(buffer,"oui")!=0)
                        return 0;
                buf.st_size = SECTOR_SIZE;
        }

        fp = fopen(file,"rb");
        if (fp == NULL)
        {
                perror(file);
                return 0;
        }

        if (fread(nouveau_boot,buf.st_size,1,fp) != 1)
        {
                printf("impossible de lire %s\n",file);
                return 0;
        }

        fclose(fp);
        return 1;
}

/* combine l'ancien secteur de boot et le nouveau, le resultat est mis dans
        le nouveau , les zones de memoire doivent etre de taille SECTOR_SIZE
*/

void combine_code(byte *nouveau_boot,byte *ancien_boot)
{
        int i;

        #define FIRST_BYTE 0x03
        #define LAST_BYTE  0x3d

        for (i=FIRST_BYTE;i<=LAST_BYTE;i++)
                nouveau_boot[i] = ancien_boot[i];
}

int main(int argc,char *argv[])
{
        int i;
        int result;
        struct position pos;
        byte buffer[SECTOR_SIZE];
        struct boot_sector *boot = (struct boot_sector *)buffer;
        byte nouveau_boot[SECTOR_SIZE],ancien_boot[SECTOR_SIZE];
        struct partition *part;

        /* quelques verifications */
        if (sizeof(struct ipart)!=16 || sizeof(struct mbr)!=SECTOR_SIZE)
        {
                printf("erreur de type!\n");
                return -1;
        }

        /* on regarde les arguments de la ligne de commande */
        if (argc != 2)
        {
                printf("usage: checkd boot-file\n");
                printf("boot-file est un fichier de %d octets contenant le\n",
                        SECTOR_SIZE);
                printf("code a mettre dans le secteur de boot de la partition\n");
                return -1;
        }

        if (!charger_boot(argv[1],nouveau_boot))
                return -1;

        printf("recherche des lecteurs physiques ...\n");
        pos.C = 0;
        pos.H = 0;
        pos.S = 1;
        for (i=0;i<256;i++)
        {
                result = read_sector((byte)i,buffer,&pos);
                if (result)
                        analyser_disque((byte)i,boot);
        }

        part = selectionner_partitions();
        if (part != NULL)
        {
                printf("modification du secteur de boot de l'unite 0x%x, secteur CHS = %d/%d/%d\n",
                        part->unite,part->premier_secteur.C,
                        part->premier_secteur.H,part->premier_secteur.S);
                printf("lecture de l'ancien secteur de boot ...");
                result = read_sector(part->unite,ancien_boot,
                        &part->premier_secteur);
                if (!result) printf("echec!\n");
                else
                {
                        printf("ok\n");
                        combine_code(nouveau_boot,ancien_boot);
                        printf("ecriture du secteur de boot ...");
                        result = write_sector(part->unite,nouveau_boot,
                                &part->premier_secteur);
                        if (result) printf("ok\n"); else printf("echec!\n");
               }
        }
        liberer_partitions();

        return 0;
}