import lief
from pwn import *
import mmap
from hackingscripts import util
from Crypto.Cipher import AES
from Crypto.Util import Counter
from Crypto.Util.number import bytes_to_long

file_path = "coredump.zst"
core = lief.parse(file_path)

class StackFrame:

    def __init__(self, rbp, rsp):
        self.rbp = rbp
        self.rsp = rsp
        assert self.rbp > self.rsp

    def get_memory(self, offset=0, size=None):
        size = util.nvl(size, len(self) - offset)
        return read_memory(self.rsp + offset, size)

    def __len__(self):
        return self.rbp - self.rsp

    def __repr__(self):
        return f"<StackFrame rbp={hex(self.rbp)} rsp={hex(self.rsp)} size={hex(len(self))}>"

def read_memory(addr, size):
    for segment in core.segments:
        if segment.type == lief.ELF.SEGMENT_TYPES.LOAD:
            start_address = segment.virtual_address
            end_address = start_address + segment.physical_size

            if start_address <= addr < end_address:
                with open(file_path, 'rb') as f:
                    with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mmapped_file:
                        offset = addr - start_address + segment.file_offset
                        mmapped_file.seek(offset)
                        data = mmapped_file.read(size)
                        return data

    raise Exception("Invalid address:", hex(addr))

for note in core.notes:
   if note.type_core == lief.ELF.NOTE_TYPES_CORE.PRSTATUS:
      details = note.details
      rsp = details[lief.ELF.CorePrStatus.REGISTERS.X86_64_RSP]
      rbp = details[lief.ELF.CorePrStatus.REGISTERS.X86_64_RBP]
      r13 = details[lief.ELF.CorePrStatus.REGISTERS.X86_64_R13]
      stack_frame = StackFrame(rbp, rsp)


print("[+] RSP at:", hex(rsp))
key = stack_frame.get_memory(0x10, 0x20)
print("[+] Got key:", key.hex())

heap_addr = r13
print("[+] Heap chunk at:", hex(heap_addr))
encrypted = read_memory(heap_addr, 0x30)
iv = encrypted[:16]
ct = encrypted[16:].rstrip(b"\x00")

print("[+] Got IV:", iv.hex())
print("[+] Got ct:", ct.hex())

ctr = Counter.new(128, initial_value=bytes_to_long(iv))
cipher = AES.new(key, AES.MODE_CTR, counter=ctr)
flag = cipher.decrypt(ct).decode().strip()
print("[+] Flag:", flag)