import re, subprocess, idaapi, ida_segment, ida_kernwin

# To install this, simply put it in your ida_install/loaders folder and open
# a `/proc/<pid>/mem` file!
#
# You might need to set `echo 0 > /proc/sys/kernel/yama/ptrace_scope` if you
# want to be able to dump processes depending on your system configuration.

# Check if the file is supported by our loader
def accept_file(li, filename):
    # Check if the filename is /proc/<pid>/mem, if so, we can handle it!
    mch = re.match("/proc/(\d+)/mem", filename)
    if not mch:
        return 0

    # We can handle this file!
    return {'format': f"{filename} dump", 'processor': 'metapc'}

# Load the file into the database!
def load_file(li, neflags, fmt):
    # Get the PID from the format
    pid = int(re.match("/proc/(\d+)/mem dump", fmt).group(1))

    # Ask the user about the bitness to use for segments
    bitness = ida_kernwin.ask_buttons(
        "64-bit", "32-bit", "Cancel", 0, "What bitness is this process?")
    if bitness == -1:
        # Cancelled
        return 0

    # Convert dialog selection to IDA's segment bitness values
    if bitness == 1:
        bitness = 2 # 64-bit
    elif bitness == 0:
        bitness = 1 # 32-bit

    # Open the maps file
    seg = idaapi.segment_t()
    with open(f"/proc/{pid}/maps") as fd:
        # Go through each line in the map
        for line in fd.readlines():
            # Parse the /proc/<pid>/maps line, super quality regex
            mch = re.match("([0-9a-f]+)-([0-9a-f]+) ([r-])([w-])([x-])[ps] [0-9a-f]+ [0-9a-f]+:[0-9a-f]+ \d+\s+(.*)", line)
            start = int(mch.group(1), 16)
            end   = int(mch.group(2), 16)
            r     = mch.group(3) == "r"
            w     = mch.group(4) == "w"
            x     = mch.group(5) == "x"
            name  = mch.group(6)

            # If this segment is usable in any way, add it to the database
            if r or w or x:
                # Mark things as code if they're executable
                seg.start_ea = start
                seg.end_ea   = end
                seg.bitness  = bitness

                # Set up permissions
                seg.perm = 0
                if r:
                    seg.perm |= ida_segment.SEGPERM_READ
                if w:
                    seg.perm |= ida_segment.SEGPERM_WRITE
                if x:
                    seg.perm |= ida_segment.SEGPERM_EXEC

                # Add the segment!
                if x:
                    idaapi.add_segm_ex(seg, name, "CODE", 0)
                else:
                    idaapi.add_segm_ex(seg, name, "DATA", 0)

                # IDA's API doesn't like seeking negative
                if start >= 0x8000000000000000:
                    print(f"Unsupported address range {start:x}-{end:x} "
                           "leaving uninitialized")
                    continue

                # Seek to the data and read it
                li.seek(start)
                data = li.read(end - start)

                # It's possible we failed to read certain areas, so only
                # `put_bytes` if we actually got a valid result
                if data:
                    # Write in the bytes
                    idaapi.put_bytes(start, data)

    # Loaded!
    return 1