/*
 *    exe2ico.c: Herramienta para sacar iconos de ejecutables y DLL de
 *    Windows para Linux.
 * Copyright (c) BatchDrake 2007 <BatchDrake@gmail.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; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 * ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 */

#include <stdio.h>
#include <libgen.h>
#include <stdlib.h>
#include <time.h>
#include <errno.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>

#define DOS_SIGNATURE 0x5a4d /* "MZ" */
#define IMAGE_NT_SIGNATURE 0x4550 /* "PE" */
#define OPTIONAL_HEADER_SIGNATURE 0x10b

#define __packed  __attribute__ ((packed))

#define _u64 unsigned long long int
#define _u32 unsigned int
#define _u16 unsigned short
#define _u8  unsigned char

#define _s64 long long int
#define _s32 int
#define _s16 short
#define _s8  char

#define RT_ICON 3

#define UNSIG(x) ( (x) & ~0x80000000)
#define error(fmt, arg...) fprintf (stderr, "error: %s: " fmt, name, ##arg)

/* Cabecera DOS del ejecutable */
typedef struct _doshdr
{
    _u16 dh_signature __packed;
    _u8  dh_dos_data[22];
    _u32 dh_stub_offset;
    _u8  dh_dos_data2[32];
    _u32 dh_pe_offset;
}
dos_header_t;

/* Cabecera PE (COFF) del ejecutable */
typedef struct _pehdr
{
    _u32 pe_signature;
    _u16 pe_arch;
    _u16 pe_sect_count;
    _u32 pe_assembly_time;
    _u32 pe_sym_table_offset;
    _u32 pe_sym_count;
    _u16 pe_ohdr_size __packed;
    _u16 pe_charact __packed;
    
}
pe_header_t;

/* Cabecera opcional / NT */
typedef struct _ohdr
{
    _u16 op_magic;
    _u16 op_linker_version;
    _u32 op_size_of_code;
    _u32 op_size_of_data;
    _u32 op_size_of_bss;
    _u32 op_entrypoint;
    _u32 op_codebase;
    _u32 op_database;
    _u32 op_imagebase;
    _u32 op_s_align;
    _u32 op_f_align;
    _u16 op_os_major;
    _u16 op_os_minor;
    _u16 op_image_major;
    _u16 op_image_minor;
    _u16 op_subsystem_major;
    _u16 op_subsystem_minor;
    _u32 op_win32_version;
    _u32 op_image_size;
    _u32 op_headers_size;
    _u32 op_checksum;
    _u16 op_nt_subsystem;
    _u16 op_dll_chars;
    _u32 op_stack_size;
    _u32 op_stack_commit;
    _u32 op_heap_size;
    _u32 op_heap_commit;
    _u32 op_loader_flags;
    _u32 op_rva_count;
}
optional_header_t;

/* Entrada del directorio de datos */
typedef struct _rva
{
    _u32 rv_rva;
    _u32 rv_size;
}
rva_directory_entry_t;

/* Cabecera de sección. La usamos para saber dónde está .rsrc */
typedef struct _shdr
{
    _u8  sh_name[8];
    union
    {
        _u32 sh_physical;
        _u32 sh_virtualsiz;
    }
    sh_addrdata;
    _u32 sh_virtual_addr;
    _u32 sh_raw_data_size;
    _u32 sh_offset;
    _u32 sh_rel_offset;
    _u32 sh_lns_offset;
    _u16 sh_rel_count;
    _u16 sh_lns_count;
    _u32 sh_chars;
}
section_header_t;

/* Cabecera del directorio de recursos */
typedef struct _imrdir
{
    _u32  rd_char;
    _u32  rd_timestamp;
    _u16  rd_major;
    _u16  rd_minor;
    _u16  rd_named;
    _u16  rd_id;
}
image_resource_directory_t; 

/* Entrada en el directorio de recursos */
typedef struct _imrdirent
{
    _u32   rd_name;
    _u32   rd_off; /* Esto es un offset relativo, y mide desplazamientos
                      desde el inicio de la sección. */
}
image_resource_directory_entry_t;

/* Información sobre el icono en el recurso */
typedef struct _ihdrinf
{
    _u32 ih_address; /* OJO: Esto es una RVA sin la dirección base del ejecutable */
    _u32 ih_size;
    _u32 ih_pad[2];
}
icon_header_info_t;

/* Estas estructuras ya se corresponden con los ficheros .ICO */
/* Cabecera del icono */
typedef struct _icohdr
{
    _u16 id_reserved;
    _u16 id_type;
    _u16 id_count;
}
icon_header_t;

/* Entrada dentro del icono, con información sobre el tamaño del mismo */
typedef struct _icodirent
{
    _u8  ie_width;
    _u8  ie_height;
    _u8  ie_colors;
    _u8  ie_reserved;
    _u16 ie_planes;
    _u16 ie_bits;
    _u32 ie_rsrc_size;
    _u32 ie_offset;
}
icon_entry_t;

/* Cabecera DIB (datos del icono per-se) */
typedef struct _bmpinfohdr
{
    _u32  dh_size; 
    _u32  dh_width; 
    _u32  dh_height; 
    _u16  dh_planes; 
    _u16  dh_bits; 
    _u32  dh_compression; 
    _u32  dh_size_image; 
    _u32  dh_x_ppm; 
    _u32  dh_y_ppm; 
    _u32  dh_clr_used; 
    _u32  dh_clr_important; 
}
dib_header_t;

/* Algunas cabeceras globales */
optional_header_t opt_header;
rva_directory_entry_t rva_headers[16];

/* Algunas variables necesarias para el comportamiento del programa */
char *name;
int   do_count;
int   dumpcnt;
int   dumptotal;

int   specific_width = -1;
int   specific_depth = -1;
char *specific_output;

/* Esto nos carga una estructura en cualquier parte del ejecutable sin
   alterar el cursor del fichero */
int load_struct (FILE *fp, _u32 off, void *ptr, _u32 size)
{
    _u32 last_seek;
    int r;
    
    last_seek = ftell (fp);
    fseek (fp, off, SEEK_SET);
    
    r = fread (ptr, size, 1, fp);
    
    fseek (fp, last_seek, SEEK_SET);
    
    return r != 1 ? -1 : 0;
}

/* Las siguientes tres funciones son lo más complicado del programa. Son las 
  que se encargan de leer el árbol de recursos. La cosa es así:

                        .rsrc                                   
                          |
                image_resource_directory_t: Indica cuántos recursos tenemos
                      RT_DIALOG
                      RT_CURSOR
                       RT_ICON: Esto nos interesa. Offset relativo + offset
                          |     de la propia sección.
                          |     <--- offset@image_resource_directory_entry_t
                          V
                image_resource_directory_t: Indica cuántos recursos de icono
                       Recurso 1            tenemos.
                       Recurso 2
                       Recurso 3
                          |
                          |     <--- offset@image_resource_directory_entry_t
                          V
                image_resource_directory_t: Indica cuántos iconos hay dentro
                       Icono 1              de éste recurso.
                       Icono 2
                       Icono 3
                          |
                          |     <--- offset@image_resource_directory_entry_t
                          V
                  icon_header_info_t
                          |
                          V     <--- Offset: RVA sin Imagebase
        *** Datos del icono (sin cabecera .ICO) ***

Como véis, es un peñazo leer los recursos de un ejecutable de Windows. Tenemos
que leer TRES directorios de recursos en árbol y luego una cabecera de icono
de la cual tenemos que despejar la maldita dirección física a partir de una
dirección virtual sin el ImageBase.
*/
                
/* Esta vuelca el icono a partir del recurso de icono */

int dump_this_icon (FILE *fp, _u32 offset, _u32 roff)
{
    image_resource_directory_t dir;
    image_resource_directory_entry_t ent;
    icon_header_info_t hdr;
    
    /* Cabeceras del .ICO */
    icon_header_t ico_hdr;
    icon_entry_t ico_ent;
    dib_header_t *dib_hdr;
    
    char *buff;
    int i;
    FILE *sfp;
    
    char file[256];
    
    if (specific_output && dumpcnt)
        return 0;
    
    load_struct (fp, offset + roff, &dir, sizeof (dir));
    
    for (i = 0; i < dir.rd_named + dir.rd_id; i++)
    {
        load_struct (fp, offset + roff + sizeof (dir) + sizeof (ent) * i, &ent, sizeof (ent));
        load_struct (fp, UNSIG (ent.rd_off) + offset, &hdr, sizeof (hdr));
    
        buff = malloc (hdr.ih_size);
        if (buff == NULL)
        {
            error ("memoria insuficiente para volcar el icono\n");
            exit (1);
        }
    
        /* Despejamos el offset sin el Imagebase */
        load_struct (fp, hdr.ih_address - rva_headers[2].rv_rva + offset, buff, hdr.ih_size);
        
        dib_hdr = (dib_header_t *) buff;
    
        if (specific_width != -1 && specific_width != dib_hdr->dh_width)
            return 0;
    
        if (specific_depth != -1 && specific_depth != dib_hdr->dh_bits)
            return 0;
    
        if (do_count)
        {
            dumpcnt++;
            return 0;
        }
    
        if (specific_output)
            strncpy (file, specific_output, 255);
        else
            sprintf (file, "%s-%04d-%d-%dx%d.ico", name, dumpcnt, ent.rd_name, dib_hdr->dh_width, (unsigned char) dib_hdr->dh_height);    
        
    
        printf ("%s [%d]: icono %d/%d (%dx%d / %d bits) ==>\n    %s... ", name, ent.rd_name, i + 1, dir.rd_named + dir.rd_id, (unsigned char) dib_hdr->dh_width, (unsigned char) dib_hdr->dh_height, dib_hdr->dh_bits, file);
        fflush (stdout);

        sfp = fopen (file, "wb");

        if (sfp == NULL)
            perror (file);
        else
        {
            /* Inicializamos el icono */
            ico_hdr.id_reserved = 0;
            ico_hdr.id_type = 1;
            ico_hdr.id_count = 1;
            fwrite (&ico_hdr, sizeof (ico_hdr), 1, sfp);
            
            /* Inicializamos la entrada */
            ico_ent.ie_width = ico_ent.ie_height = dib_hdr->dh_width;
            ico_ent.ie_colors = 0;
            ico_ent.ie_reserved = 0;
            ico_ent.ie_planes = 0;
            ico_ent.ie_bits = 0;
            ico_ent.ie_rsrc_size = hdr.ih_size;
            ico_ent.ie_offset = sizeof (icon_header_t) + sizeof (icon_entry_t);
            
            fwrite (&ico_ent, sizeof (ico_ent), 1, sfp);
        
            /* Escribimos todos los datos del icono */
            fwrite (buff, hdr.ih_size, 1, sfp);
            fclose (sfp);
            dumpcnt++;
            dumptotal++;
        
            printf ("hecho (%d bytes escritos)\n\n", hdr.ih_size, sizeof (icon_header_t) + sizeof (icon_entry_t));
        }
    
        free (buff);
        
    }
    
    return 0;
}

/* Esto busca recursos de icono en un RT_ICON */
int analyze_icons_in_rsrc (FILE *fp, _u32 offset, _u32 roff)
{
    image_resource_directory_t dir;
    image_resource_directory_entry_t ent;
    int i;

    
    load_struct (fp, offset + roff, &dir, sizeof (dir));
    
    for (i = 0; i < dir.rd_named + dir.rd_id; i++)
    {
        load_struct (fp, offset + roff + sizeof (dir) + sizeof (ent) * i, &ent, sizeof (ent));
        dump_this_icon (fp, offset, UNSIG (ent.rd_off));
    }
    
    return 0;
}

/* Esto busca el recurso RT_ICON en .rsrc */
int analyze_rsrc_section (FILE *fp, _u32 offset, _u32 size)
{
    image_resource_directory_t dir;
    image_resource_directory_entry_t ent;

    int i;
    
    load_struct (fp, offset, &dir, sizeof (dir)); 
    
    for (i = 0; i < dir.rd_named + dir.rd_id; i++)
    {
        load_struct (fp, offset + sizeof (dir) + sizeof (ent) * i, &ent, sizeof (ent));
        if (ent.rd_name == RT_ICON)
            analyze_icons_in_rsrc (fp, offset, UNSIG (ent.rd_off));
        /* VirtualOffsetYouHave-ImageBase-VirtualOffsetOfSection+RawOffsetOfSection */
        
    }
    
    return 0;
}

void parse_exe_file (char *file)
{
    FILE *fp;
    dos_header_t dos_header;
    pe_header_t  pe_header;
    section_header_t      sect_header;
    int i;
    
    dumpcnt = 0;
    fp = fopen (file, "rb");
    name = basename (file);
    
    if (fp == NULL)
    {
        perror (file);
        return;
    }
    
    fread (&dos_header, sizeof (dos_header_t), 1, fp);
    if (dos_header.dh_signature != DOS_SIGNATURE)
    {
        error ("no es un ejecutable MZ válido\n");
        fclose (fp);
        return;
    }
    
    if (!dos_header.dh_pe_offset)
    {
        error ("este ejecutable es de MS-DOS. Los ejecutables de MS-DOS no tienen iconos\n");
        fclose (fp);
        return;
    }
    
    fseek (fp, dos_header.dh_pe_offset, SEEK_SET);
    
    fread (&pe_header, sizeof (pe_header_t), 1, fp);
    
    if (pe_header.pe_signature != IMAGE_NT_SIGNATURE)
    {
        error ("esto no parece ser un ejecutable PE válido\n");
        error ("IMAGE_NT_SIGNATURE inesperado (0x%08x)\n", pe_header.pe_signature);
        fclose (fp);
        return;
    }

    if (pe_header.pe_ohdr_size != 224)
    {
        error ("la cabecera opcional del ejecutable no coincide con la esperada\n", pe_header.pe_ohdr_size);
        fclose (fp);
        return;
    }
    
    fread (&opt_header, sizeof (optional_header_t), 1, fp);
    if (opt_header.op_magic != OPTIONAL_HEADER_SIGNATURE)
    {
        error ("este ejecutable no tiene una cabecera opcional válida\n");
        fclose (fp);
        return;
    }
    
    
    if (opt_header.op_rva_count > 16)
    {
        error ("el directorio de datos tiene demasiadas entradas\n");
        fclose (fp);
        return;
    }
    
    for (i = 0; i < opt_header.op_rva_count; i++)
        fread (&rva_headers[i], sizeof (rva_directory_entry_t), 1, fp);
    
    if (pe_header.pe_sect_count)
    {
        for (i = 0; i < pe_header.pe_sect_count; i++)
        {
            fread (&sect_header, sizeof (section_header_t), 1, fp);
            
            if (rva_headers[2].rv_rva == sect_header.sh_virtual_addr)
                analyze_rsrc_section (fp, sect_header.sh_offset, sect_header.sh_raw_data_size);
        }
    }
    
    if (!dumpcnt)
    {
        if (specific_width != -1 || specific_depth != -1)
            fprintf (stderr, "%s: no hay iconos que sigan el criterio\n", name);
        else
            fprintf (stderr, "%s: el ejecutable no tiene iconos\n", name);
    }
    else
    {
        printf ("%s: %d iconos", name, dumpcnt);
    
        if (do_count)
            printf ("\n");
        else
            printf (", %d volcados\n", dumpcnt);
    }
    
    fclose (fp);
}

void help (char *argv0)
{
    fprintf (stderr, "Herramienta de extración de iconos de ejecutables de Windows\n");
    fprintf (stderr, "Por Gonzalo J. Carracedo <BatchDrake@gmail.com>\n\n");
    fprintf (stderr, "FORMA DE USO:\n");
    fprintf (stderr, "    %s [opciones] fichero.exe [fichero2.exe [fichero3.dll [fichero4.scr]]]\n", argv0);
    fprintf (stderr, "\nOPCIONES:\n");
    fprintf (stderr, "    -c, --count:             Sólo cuenta los iconos (no extrae)\n");
    fprintf (stderr, "    -w, --width <tamaño>:    Extrae sólo aquellos iconos de ese ancho\n");
    fprintf (stderr, "    -d, --depth <colores>:   Extrae aquellos iconos que tengan esa profundidad\n");
    fprintf (stderr, "    -o, --output <archivo>:  Extrae el primer icono a <archivo>\n");
    fprintf (stderr, "    --help:                  Esta ayuda\n");
    fprintf (stderr, "\nEste programa es libre y se distribuye bajo los términos de la licencia\n");
    fprintf (stderr, "GNU/GPL versión 2.\n");
}

int main (int argc, char **argv)
{
    FILE *fp;
    int i;
    int mul;
    
    mul = 0;
    
    if (argc < 2)
    {
        fprintf (stderr, "%s: Se speraba un argumento\n", argv[0]);
        help (argv[0]);
        exit (1);
    }
    
    
    for (i = 1; i < argc; i++)
    {
        if (strcmp (argv[i], "-c") == 0 || strcmp (argv[i], "--count") == 0)
        {
            do_count = 1;
            continue;
        }
        else if (strcmp (argv[i], "--help") == 0)
        {
            help (argv[0]);
            exit (0);
        }
        else if (strcmp (argv[i], "-w") == 0 || strcmp (argv[i], "--width") == 0)
        {
            if (++i >= argc)
            {
                fprintf (stderr, "%s: La opción %s espera un valor\n\n", argv[0], argv[i - 1]);
                help (argv[0]);
                exit (1);
            }
        
            if (!sscanf (argv[i], "%u", &specific_width))
            {
                fprintf (stderr, "%s: La opción %s espera un valor numérico\n\n", argv[0], argv[i - 1]);
                help (argv[0]);
                exit (1);
            }
        
            continue;
        }
        else if (strcmp (argv[i], "-d") == 0 || strcmp (argv[i], "--depth") == 0)
        {
            if (++i >= argc)
            {
                fprintf (stderr, "%s: La opción %s espera un valor\n\n", argv[0], argv[i - 1]);
                help (argv[0]);
                exit (1);
            }
        
            if (!sscanf (argv[i], "%u", &specific_depth))
            {
                fprintf (stderr, "%s: La opción %s espera un valor numérico\n\n", argv[0], argv[i - 1]);
                help (argv[0]);
                exit (1);
            }
        
            continue;
        }
        else if (strcmp (argv[i], "-o") == 0 || strcmp (argv[i], "--output") == 0)
        {
            if (++i >= argc)
            {
                fprintf (stderr, "%s: La opción %s espera un valor\n\n", argv[0], argv[i - 1]);
                help (argv[0]);
                exit (1);
            }
        
            specific_output = argv[i];
        
            continue;
        }
    
        parse_exe_file (argv[i]);
        mul++;
    }
    
    if (!mul)
    {
        fprintf (stderr, "%s: Se speraba un fichero\n", argv[0]);
        help (argv[0]);
        exit (1);
    }
    
    if (mul > 2)
        printf ("Total: %d iconos volcados en todos los archivos\n", dumptotal);

    return 0;
}