1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
//! Implementation of the entrypoint instruction in yara rules.
//!
//! This depends on the `object` feature.

use object::{
    coff::SectionTable,
    pe::{
        ImageDosHeader, ImageNtHeaders32, ImageNtHeaders64, IMAGE_FILE_MACHINE_AMD64,
        IMAGE_FILE_MACHINE_I386,
    },
    read::{
        elf::{ElfFile, ElfFile32, ElfFile64, FileHeader},
        pe::{ImageNtHeaders, ImageOptionalHeader},
    },
    Endianness, FileKind, LittleEndian as LE,
};

use super::Value;
use crate::module::elf;

pub(super) fn get_pe_or_elf_entry_point(mem: &[u8]) -> Option<Value> {
    match FileKind::parse(mem).ok()? {
        FileKind::Pe32 => parse_pe::<ImageNtHeaders32>(mem),
        FileKind::Pe64 => parse_pe::<ImageNtHeaders64>(mem),
        FileKind::Elf32 => parse_elf(&ElfFile32::parse(mem).ok()?, mem),
        FileKind::Elf64 => parse_elf(&ElfFile64::parse(mem).ok()?, mem),
        _ => None,
    }
}

fn parse_pe<Pe: ImageNtHeaders>(mem: &[u8]) -> Option<Value> {
    let dos_header = ImageDosHeader::parse(mem).ok()?;
    let mut offset = dos_header.nt_headers_offset().into();
    let (nt_headers, _) = Pe::parse(mem, &mut offset).ok()?;
    let opt_hdr = nt_headers.optional_header();
    let sections = nt_headers.sections(mem, offset).ok()?;

    // For some reasons, those tests exists here, but not in pe module...
    let machine = nt_headers.file_header().machine.get(LE);
    if machine != IMAGE_FILE_MACHINE_I386 && machine != IMAGE_FILE_MACHINE_AMD64 {
        return None;
    }

    let ep = opt_hdr.address_of_entry_point();

    Some(Value::Integer(
        pe_rva_to_file_offset(&sections, ep).unwrap_or(0),
    ))
}

// This reimplements the `yr_pe_rva_to_offset` function from libyara.
//
// This is not the same as the function implemented in the pe module and it has it own set
// of particularities...
fn pe_rva_to_file_offset(sections: &SectionTable, va: u32) -> Option<i64> {
    let mut nearest_section_va = 0;
    let mut nearest_section_offset = 0;
    for section in sections.iter().take(60) {
        let section_va = section.virtual_address.get(LE);
        if va >= section_va && nearest_section_va <= section_va {
            nearest_section_va = section_va;
            nearest_section_offset = section.pointer_to_raw_data.get(LE);
        }
    }

    i64::from(nearest_section_offset).checked_add(i64::from(va - nearest_section_va))
}

fn parse_elf<Elf: FileHeader<Endian = Endianness>>(
    file: &ElfFile<Elf>,
    mem: &[u8],
) -> Option<Value> {
    let ep = elf::entry_point(file, mem)?;
    match i64::try_from(ep) {
        Ok(v) => Some(Value::Integer(v)),
        Err(_) => None,
    }
}