#!/usr/bin/env python3
import struct
from typing import Final

def build_message() -> bytes:
    """
    Build the data segment containing the message.
    """
    return b"Hello, world\n"  # 13 bytes

def build_code(message_disp: int) -> bytes:
    """
    Build the machine-code segment that performs:
      - write syscall to print the message, and
      - exit syscall to terminate the program.

    The instruction 'lea rsi, [rip+disp]' will use the provided
    displacement (in bytes) to compute the address of the message.
    """
    # Assemble the instructions:
    #   mov rax, 1              ; syscall number for write
    #   mov rdi, 1              ; file descriptor (stdout)
    #   lea rsi, [rip+disp]     ; pointer to the message (disp provided)
    #   mov rdx, 13             ; length of the message
    #   syscall                 ; invoke write
    #   mov rax, 60             ; syscall number for exit
    #   xor rdi, rdi            ; exit code 0
    #   syscall                 ; invoke exit
    code_bytes: Final[bytes] = bytes([
        # mov rax, 1
        0x48, 0xc7, 0xc0, 0x01, 0x00, 0x00, 0x00,
        # mov rdi, 1
        0x48, 0xc7, 0xc7, 0x01, 0x00, 0x00, 0x00,
        # lea rsi, [rip+disp]
        0x48, 0x8d, 0x35,
        # disp (4 bytes placeholder, will be replaced below)
        0x00, 0x00, 0x00, 0x00,
        # mov rdx, 13
        0x48, 0xc7, 0xc2, 0x0d, 0x00, 0x00, 0x00,
        # syscall (write)
        0x0f, 0x05,
        # mov rax, 60
        0x48, 0xc7, 0xc0, 0x3c, 0x00, 0x00, 0x00,
        # xor rdi, rdi
        0x48, 0x31, 0xff,
        # syscall (exit)
        0x0f, 0x05
    ])
    # The lea instruction starts at offset 14 and occupies 7 bytes.
    # The displacement is stored in 4 bytes at offset 17.
    # The effective address for the message will be computed relative
    # to the address immediately after the lea instruction (offset 21).
    disp_bytes: bytes = struct.pack("<i", message_disp)
    code_array = bytearray(code_bytes)
    code_array[17:21] = disp_bytes
    return bytes(code_array)

def build_segment() -> bytes:
    """
    Construct the full segment (code followed by data).

    The correct displacement in the lea instruction must be the difference
    between the start of the data and the address following the lea instruction.
    """
    message: bytes = build_message()
    # First, build the code with a dummy displacement.
    dummy_code = build_code(0)
    # The lea instruction's effective address is computed from offset 21.
    # The correct displacement is:
    #     len(code) + 0 (start of message) - 21 = len(code) - 21
    correct_disp: int = len(dummy_code) - 21
    code: bytes = build_code(correct_disp)
    return code + message

def build_elf_header(segment_offset: int, entry_point: int, ph_offset: int, total_size: int) -> bytes:
    """
    Build the ELF header (64 bytes) with the given parameters.

    - segment_offset: Offset in the file where the code+data segment begins.
    - entry_point: Virtual address at which execution begins.
    - ph_offset: Offset to the program header table (immediately after ELF header).
    - total_size: Total file size (used for segment size fields).
    """
    # e_ident: magic number, 64-bit, little-endian, version 1, etc.
    e_ident: bytes = b'\x7fELF' + bytes([2, 1, 1, 0]) + bytes(8)
    e_type: int = 2            # Executable file (ET_EXEC)
    e_machine: int = 0x3e        # x86-64 (EM_X86_64)
    e_version: int = 1
    e_entry: int = entry_point  # Entry point virtual address
    e_phoff: int = ph_offset    # Program header table offset
    e_shoff: int = 0            # No section header table in this minimal binary
    e_flags: int = 0
    e_ehsize: int = 64          # ELF header size
    e_phentsize: int = 56       # Program header entry size for ELF64 is 56 bytes
    e_phnum: int = 1            # One program header entry
    e_shentsize: int = 0        # Not used
    e_shnum: int = 0            # Not used
    e_shstrndx: int = 0         # Not used

    return struct.pack(
        "<16sHHIQQQIHHHHHH",
        e_ident,
        e_type,
        e_machine,
        e_version,
        e_entry,
        e_phoff,
        e_shoff,
        e_flags,
        e_ehsize,
        e_phentsize,
        e_phnum,
        e_shentsize,
        e_shnum,
        e_shstrndx
    )

def build_program_header(total_size: int, base_vaddr: int) -> bytes:
    """
    Build the program header (PT_LOAD) that maps the entire file into memory.

    - total_size: Total file size (headers + segment).
    - base_vaddr: Base virtual address at which the segment is loaded.
    """
    p_type: int = 1          # PT_LOAD
    p_flags: int = 5         # PF_R | PF_X (readable and executable)
    p_offset: int = 0        # Segment starts at file offset 0
    p_vaddr: int = base_vaddr   # Virtual address where the segment is loaded
    p_paddr: int = base_vaddr   # Physical address (unused on Linux)
    p_filesz: int = total_size
    p_memsz: int = total_size
    p_align: int = 0x200000    # Typical page alignment

    return struct.pack(
        "<IIQQQQQQ",
        p_type,
        p_flags,
        p_offset,
        p_vaddr,
        p_paddr,
        p_filesz,
        p_memsz,
        p_align
    )

def create_elf() -> bytes:
    """
    Assemble the final ELF binary by combining:
      - the ELF header,
      - the program header, and
      - the segment (code + data).
    """
    base_vaddr: Final[int] = 0x400000
    elf_header_size: int = 64
    ph_size: int = 56
    segment_offset: int = elf_header_size + ph_size
    segment: bytes = build_segment()
    total_size: int = segment_offset + len(segment)
    entry_point: int = base_vaddr + segment_offset

    elf_hdr: bytes = build_elf_header(segment_offset, entry_point, elf_header_size, total_size)
    prog_hdr: bytes = build_program_header(total_size, base_vaddr)
    return elf_hdr + prog_hdr + segment

def write_elf(file_name: str, elf_binary: bytes) -> None:
    """
    Write the ELF binary to a file.
    """
    with open(file_name, "wb") as f:
        f.write(elf_binary)

def main() -> None:
    elf_binary: bytes = create_elf()
    file_name: str = "hello"
    write_elf(file_name, elf_binary)
    print(f"ELF binary written to {file_name}.")
    print("Mark it executable (e.g., 'chmod +x hello') and run it on a Linux x86_64 system.")

if __name__ == "__main__":
    main()