import sys

import os
import mmap

def main():
if len(sys.argv) != 3:
print("USAGE: python fixPng.py [directory] [output]")
else:
decode_png()

def decode_png():
old_png_bytes: bytes = b"isnvsedq"
new_png_bytes: bytes = b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a"
old_jpg_bytes: bytes = b"\x69"
new_jpg_bytes: bytes = b"\xff\xd8\xff"

path: str = sys.argv[1]

output: str = sys.argv[2]

if not os.path.isdir(path):
raise RuntimeError("Passed Directory is not a directory!")

if not os.path.exists(output):
os.makedirs(output)

all_files: list[str] = os.listdir(path)

for file_name in all_files:

old_file_path: str = f"{path}/{file_name}"

if os.path.isdir(old_file_path) or not os.path.isfile(old_file_path):

continue

with open(old_file_path, "r+b") as original_file:

memory_mapped_file = mmap.mmap(original_file.fileno(), 0)

if memory_mapped_file[:8] == old_png_bytes:
with open(f"{output}/{file_name}.png", "wb") as new_file:
print("png")
new_file.write(new_png_bytes)
new_file.write(memory_mapped_file[8:])

elif memory_mapped_file[:1] == old_jpg_bytes:

with open(f"{output}/{file_name}.jpg", "wb") as new_file:
print("jpg")
new_file.write(new_jpg_bytes)
new_file.write(memory_mapped_file[3:])

def old_method():
import sys
import os
import mmap

if __name__ == "__main__":
if len(sys.argv) != 3:
print("USAGE: python fixPng.py [directory] [output]")
 else:
path = sys.argv[1]
output = sys.argv[2]
if not os.path.exists(output):
os.makedirs(output)
if os.path.isdir(path):
allFilesArray = current_files = os.listdir(path)
for file in allFilesArray:
if not os.path.isdir(f"{path}/{file}"):
filename = file
print(filename)
with open(path + "/" + filename, "r+b") as oriFile:
mm = mmap.mmap(oriFile.fileno(), 0)
if mm[:3] == b"enc":
with open(output + "/" + filename + ".png", "wb")
as newFile:
newFile.write(b"\x89\x50\x4e\x47\x0d\x0a\x1a\
x0a")
newFile.write(mm[8:])
elif mm[:1] == b"\x69":
with open(output + "/" + filename + ".jpg", "wb")
as newFile:
newFile.write(b"\xff\xd8\xff")
newFile.write(mm[3:])

if __name__ == "__main__":
old_method()