Adds useful metadata and color information to PNG screenshots. Meant to run right after it was taken.

screenshot_add_metadata.py

import struct
import sys
import os
import re
import zlib
from datetime import datetime, timezone
import ctypes
import ctypes.wintypes as wintypes

# Load DLLs
user32 = ctypes.WinDLL("user32", use_last_error=True)
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
dwmapi = ctypes.WinDLL("dwmapi")

PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
DWMWA_CLOAKED = 14
QDC_ONLY_ACTIVE_PATHS = 0x00000002
DISPLAYCONFIG_DEVICE_INFO_GET_SOURCE_NAME = 1
DISPLAYCONFIG_DEVICE_INFO_GET_ADVANCED_COLOR_INFO = 9
MONITOR_DEFAULTTONEAREST = 0x00000002

class GUITHREADINFO(ctypes.Structure):
	_fields_ = [
		("cbSize", wintypes.DWORD),
		("flags", wintypes.DWORD),
		("hwndActive", wintypes.HWND),
		("hwndFocus", wintypes.HWND),
		("hwndCapture", wintypes.HWND),
		("hwndMenuOwner", wintypes.HWND),
		("hwndMoveSize", wintypes.HWND),
		("hwndCaret", wintypes.HWND),
		("rcCaret", wintypes.RECT)
	]

class MONITORINFOEX(ctypes.Structure):
	_fields_ = [
		("cbSize", wintypes.DWORD),
		("rcMonitor", wintypes.RECT),
		("rcWork", wintypes.RECT),
		("dwFlags", wintypes.DWORD),
		("szDevice", wintypes.WCHAR * 32),
	]

class LUID(ctypes.Structure):
	_fields_ = [("LowPart", wintypes.DWORD), ("HighPart", wintypes.LONG)]


class DISPLAYCONFIG_RATIONAL(ctypes.Structure):
	_fields_ = [("Numerator", ctypes.c_uint32), ("Denominator", ctypes.c_uint32)]


class DISPLAYCONFIG_PATH_SOURCE_INFO(ctypes.Structure):
	_fields_ = [
		("adapterId", LUID),
		("id", ctypes.c_uint32),
		("modeInfoIdx", ctypes.c_uint32),
		("statusFlags", ctypes.c_uint32),
	]


class DISPLAYCONFIG_PATH_TARGET_INFO(ctypes.Structure):
	_fields_ = [
		("adapterId", LUID),
		("id", ctypes.c_uint32),
		("modeInfoIdx", ctypes.c_uint32),
		("outputTechnology", ctypes.c_uint32),
		("rotation", ctypes.c_uint32),
		("scaling", ctypes.c_uint32),
		("refreshRate", DISPLAYCONFIG_RATIONAL),
		("scanLineOrdering", ctypes.c_uint32),
		("targetAvailable", wintypes.BOOL),
		("statusFlags", ctypes.c_uint32),
	]


class DISPLAYCONFIG_PATH_INFO(ctypes.Structure):
	_fields_ = [
		("sourceInfo", DISPLAYCONFIG_PATH_SOURCE_INFO),
		("targetInfo", DISPLAYCONFIG_PATH_TARGET_INFO),
		("flags", ctypes.c_uint32),
	]


class DISPLAYCONFIG_MODE_INFO(ctypes.Structure):
	# Union of target/source/desktop modes; largest member is 48 bytes, header is 16
	_fields_ = [("_opaque", ctypes.c_byte * 64)]


class DISPLAYCONFIG_DEVICE_INFO_HEADER(ctypes.Structure):
	_fields_ = [
		("type", ctypes.c_uint32),
		("size", ctypes.c_uint32),
		("adapterId", LUID),
		("id", ctypes.c_uint32),
	]


class DISPLAYCONFIG_ADVANCED_COLOR_INFO(ctypes.Structure):
	_fields_ = [
		("header", DISPLAYCONFIG_DEVICE_INFO_HEADER),
		("value", ctypes.c_uint32),   # bitfield: bit0=supported, bit1=enabled
		("colorEncoding", ctypes.c_uint32),
		("bitsPerColorChannel", ctypes.c_uint32),
	]


class DISPLAYCONFIG_SOURCE_DEVICE_NAME(ctypes.Structure):
	_fields_ = [
		("header", DISPLAYCONFIG_DEVICE_INFO_HEADER),
		("viewGdiDeviceName", wintypes.WCHAR * 32),
	]

def is_window_cloaked(hwnd):
	cloaked = ctypes.c_int(0)
	dwmapi.DwmGetWindowAttribute(hwnd, DWMWA_CLOAKED, ctypes.byref(cloaked), ctypes.sizeof(cloaked))
	return cloaked.value != 0

def get_process_name_from_pid(pid):
	h_process = kernel32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid)
	if not h_process:
		return None
	try:
		exe_path = ctypes.create_unicode_buffer(260)
		size = wintypes.DWORD(260)
		if kernel32.QueryFullProcessImageNameW(h_process, 0, exe_path, ctypes.byref(size)):
			return os.path.basename(exe_path.value)
	finally:
		kernel32.CloseHandle(h_process)
	return None

def get_active_process_name():
	hwnd = user32.GetForegroundWindow()
	if not hwnd or not user32.IsWindowVisible(hwnd) or is_window_cloaked(hwnd):
		return None

	length = user32.GetWindowTextLengthW(hwnd)
	if length == 0:
		return None

	class_name = ctypes.create_unicode_buffer(256)
	user32.GetClassNameW(hwnd, class_name, 256)

	# Progman/WorkerW = Wallpaper/Icons. Shell_TrayWnd = Taskbar.
	if class_name.value in ["Progman", "WorkerW", "Shell_TrayWnd"]:
		return None

	pid = wintypes.DWORD()
	user32.GetWindowThreadProcessId(hwnd, ctypes.byref(pid))

	name = get_process_name_from_pid(pid.value)

	if not name:
		return None

	# UWP Handling
	if name and name.lower() == "applicationframehost.exe":
		gui_info = GUITHREADINFO()
		gui_info.cbSize = ctypes.sizeof(GUITHREADINFO)
		user32.GetGUIThreadInfo(0, ctypes.byref(gui_info))

		if gui_info.hwndFocus:
			child_pid = wintypes.DWORD()
			user32.GetWindowThreadProcessId(gui_info.hwndFocus, ctypes.byref(child_pid))
			uwp_name = get_process_name_from_pid(child_pid.value)
			if uwp_name:
				name = uwp_name

	# Don't care about desktop
	if name.lower() == "explorer.exe":
		if class_name.value == "CabinetWClass":
			return "Explorer"
		return None

	return os.path.splitext(name)[0]

def is_hdr_enabled():
	user32 = ctypes.windll.user32

	# Resolve active monitor device name (e.g. "\\.\DISPLAY2")
	hwnd = user32.GetForegroundWindow()
	hmonitor = user32.MonitorFromWindow(hwnd, MONITOR_DEFAULTTONEAREST)
	mon_info = MONITORINFOEX()
	mon_info.cbSize = ctypes.sizeof(MONITORINFOEX)
	if not user32.GetMonitorInfoW(hmonitor, ctypes.byref(mon_info)):
		return False
	active_device = mon_info.szDevice

	num_paths = ctypes.c_uint32(0)
	num_modes = ctypes.c_uint32(0)
	ret = user32.GetDisplayConfigBufferSizes(
		QDC_ONLY_ACTIVE_PATHS, ctypes.byref(num_paths), ctypes.byref(num_modes)
	)
	if ret != 0:
		print(f"GetDisplayConfigBufferSizes failed: {ret}")
		return False

	paths = (DISPLAYCONFIG_PATH_INFO * num_paths.value)()
	modes = (DISPLAYCONFIG_MODE_INFO * num_modes.value)()
	ret = user32.QueryDisplayConfig(
		QDC_ONLY_ACTIVE_PATHS,
		ctypes.byref(num_paths), paths,
		ctypes.byref(num_modes), modes,
		None,
	)
	if ret != 0:
		print(f"QueryDisplayConfig failed: {ret}")
		return False

	for i in range(num_paths.value):
		# Match path to active monitor via its GDI source device name
		src_name = DISPLAYCONFIG_SOURCE_DEVICE_NAME()
		src_name.header.type = DISPLAYCONFIG_DEVICE_INFO_GET_SOURCE_NAME
		src_name.header.size = ctypes.sizeof(DISPLAYCONFIG_SOURCE_DEVICE_NAME)
		src_name.header.adapterId = paths[i].sourceInfo.adapterId
		src_name.header.id = paths[i].sourceInfo.id

		if user32.DisplayConfigGetDeviceInfo(ctypes.byref(src_name.header)) != 0:
			continue
		if src_name.viewGdiDeviceName != active_device:
			continue

		target = paths[i].targetInfo
		color_info = DISPLAYCONFIG_ADVANCED_COLOR_INFO()
		color_info.header.type = DISPLAYCONFIG_DEVICE_INFO_GET_ADVANCED_COLOR_INFO
		color_info.header.size = ctypes.sizeof(DISPLAYCONFIG_ADVANCED_COLOR_INFO)
		color_info.header.adapterId = target.adapterId
		color_info.header.id = target.id

		if user32.DisplayConfigGetDeviceInfo(ctypes.byref(color_info.header)) != 0:
			return False

		enabled = bool(color_info.value & 0x2)
		supported = bool(color_info.value & 0x1)
		print(f"Display {i} ({active_device}): HDR {'enabled' if enabled else 'disabled'} (supported={supported})")
		return enabled

	return False

def get_active_icc_path():
	gdi32 = ctypes.windll.gdi32
	user32 = ctypes.windll.user32

	hwnd = user32.GetForegroundWindow()
	hmonitor = user32.MonitorFromWindow(hwnd, MONITOR_DEFAULTTONEAREST)

	info = MONITORINFOEX()
	info.cbSize = ctypes.sizeof(MONITORINFOEX)
	if not user32.GetMonitorInfoW(hmonitor, ctypes.byref(info)):
		return None

	device_name = info.szDevice
	hdc = gdi32.CreateDCW(device_name, device_name, None, None)
	if not hdc:
		return None

	try:
		size = wintypes.DWORD(0)
		gdi32.GetICMProfileW(hdc, ctypes.byref(size), None)

		if size.value == 0:
			return None

		buffer = ctypes.create_unicode_buffer(size.value)
		if gdi32.GetICMProfileW(hdc, ctypes.byref(size), buffer):
			return buffer.value
		return None
	finally:
		gdi32.DeleteDC(hdc)

def read_xyz_tag(file_handle, offset):
	file_handle.seek(offset)
	data = file_handle.read(20)
	if len(data) < 20:
		return None

	signature, _, x_val, y_val, z_val = struct.unpack('>4sIiii', data)

	if signature != b'XYZ ':
		return None

	return (x_val / 65536.0, y_val / 65536.0, z_val / 65536.0)

def calculate_xy_chromaticity(x_val, y_val, z_val):
	total = x_val + y_val + z_val
	if total == 0:
		return 0.0, 0.0
	return (x_val / total, y_val / total)

def extract_chromaticities(profile_path):
	if not profile_path or not os.path.exists(profile_path):
		return None

	chromaticities = {}

	with open(profile_path, 'rb') as file_handle:
		file_handle.seek(128)
		tag_count_data = file_handle.read(4)
		if len(tag_count_data) < 4:
			return None

		tag_count = struct.unpack('>I', tag_count_data)[0]
		tags = {}

		for _ in range(tag_count):
			tag_data = file_handle.read(12)
			if len(tag_data) < 12:
				break

			signature, offset, _ = struct.unpack('>4sII', tag_data)
			tags[signature] = offset

		target_tags = {
			b'rXYZ': 'red',
			b'gXYZ': 'green',
			b'bXYZ': 'blue',
			b'wtpt': 'white_point'
		}

		for byte_sig, string_name in target_tags.items():
			if byte_sig in tags:
				xyz_values = read_xyz_tag(file_handle, tags[byte_sig])
				if xyz_values:
					xy_values = calculate_xy_chromaticity(*xyz_values)
					chromaticities[string_name] = {
						'XYZ': xyz_values,
						'xy': xy_values
					}

	return chromaticities

# Determine captured process name
process_name = get_active_process_name()
print("Active process:", process_name)

profile_path = get_active_icc_path()
if profile_path:
	print(f"Active Profile: {profile_path}")
	chromaticity_data = extract_chromaticities(profile_path)

	if chromaticity_data:
		for color_name, data in chromaticity_data.items():
			xyz_vals = data['XYZ']
			xy_vals = data['xy']
			print(f"{color_name.capitalize()}:")
			print(f"  XYZ: ({xyz_vals[0]:.4f}, {xyz_vals[1]:.4f}, {xyz_vals[2]:.4f})")
			print(f"  xy:  ({xy_vals[0]:.4f}, {xy_vals[1]:.4f})")
	else:
		print("Could not extract chromaticity data. Ensure the profile is a matrix/TRC type.")
else:
	print("No active ICC profile found.")

hdr_enabled = is_hdr_enabled()

def to_fixed(v): return int(round(v * 100000))
# cHRM values
chrm_payload = struct.pack(">8I",
	to_fixed(chromaticity_data['white_point']['xy'][0]), to_fixed(chromaticity_data['white_point']['xy'][1]),
	to_fixed(chromaticity_data['red']['xy'][0]),   to_fixed(chromaticity_data['red']['xy'][1]),
	to_fixed(chromaticity_data['green']['xy'][0]), to_fixed(chromaticity_data['green']['xy'][1]),
	to_fixed(chromaticity_data['blue']['xy'][0]),  to_fixed(chromaticity_data['blue']['xy'][1])
)

def png_crc32(data: bytes) -> int:
	crc = 0xFFFFFFFF
	for b in data:
		crc ^= b
		for _ in range(8):
			crc = (crc >> 1) ^ (0xEDB88320 if crc & 1 else 0)
	return crc ^ 0xFFFFFFFF

def make_chunk(type_: bytes, payload: bytes) -> bytes:
	crc = png_crc32(type_ + payload)
	return struct.pack(">I", len(payload)) + type_ + payload + struct.pack(">I", crc)

def make_chrm_chunk(): return make_chunk(b"cHRM", chrm_payload)

def make_iccp_chunk(path: str) -> bytes:
	with open(path, "rb") as f:
		profile_data = f.read()
	name = os.path.splitext(os.path.basename(path))[0][:79].encode("latin-1")
	payload = name + b"\x00\x00" + zlib.compress(profile_data)
	return make_chunk(b"iCCP", payload)

def make_time_chunk(dt: datetime):
	dt = dt.astimezone(timezone.utc)
	payload = struct.pack(">HBBBBB", dt.year, dt.month, dt.day, dt.hour, dt.minute, dt.second)
	return make_chunk(b"tIME", payload)

def make_text_chunk(keyword: str, text: str):
	payload = keyword.encode("latin-1") + b"\0" + text.encode("latin-1")
	return make_chunk(b"tEXt", payload)

def insert_before_first_idat(png: bytearray, chunk: bytes):
	pos = 8
	while pos < len(png) - 12:
		length = struct.unpack_from(">I", png, pos)[0]
		ctype = png[pos+4:pos+8]
		if ctype == b"IDAT":
			png[pos:pos] = chunk
			return
		pos += 12 + length
	png[-12:-12] = chunk


# Main Logic
if len(sys.argv) < 2:
	print('Usage: python script.py "input.png"')
	sys.exit(1)

input_path = sys.argv[1].strip('"').replace("\\\\", "\\")
if not os.path.isfile(input_path):
	print(f"File not found: {input_path}")
	sys.exit(1)

norm_input = os.path.normcase(os.path.normpath(input_path))
found_obj = None

# Read PNG Data
with open(input_path, "rb") as f:
	png = bytearray(f.read())

# Strip existing metadata
pos = 8
while pos < len(png) - 12:
	length = struct.unpack_from(">I", png, pos)[0]
	ctype = png[pos+4:pos+8]
	if ctype in (b"iCCP", b"cHRM", b"tEXt", b"tIME"):
		del png[pos:pos + 12 + length]
		continue
	pos += 12 + length

# Determine DateTime
dt = None

basename = os.path.basename(input_path)
print("Parsing filename for data...")
# Supported formats: Screenshot YYYY-MM-DD HHMMSS.png
pattern = r'^Screenshot\s+(?P<year>\d{4})-(?P<mo>\d{2})-(?P<day>\d{2})\s+(?P<h>\d{2})(?P<mi>\d{2})(?P<s>\d{2})'
match = re.match(pattern, basename)
local_tz = datetime.now().astimezone().tzinfo
if match:
	year = int(match.group('year'))
	mo = int(match.group('mo'))
	day = int(match.group('day'))
	hour = int(match.group('h'))
	minute = int(match.group('mi'))
	second = int(match.group('s'))
	# Interpret the filename timestamp as local time, then convert to UTC
	dt_local = datetime(year, mo, day, hour, minute, second).replace(tzinfo=local_tz)
	dt = dt_local.astimezone(timezone.utc)
	print(f"Parsed date from filename (UTC): {dt.isoformat()}")
else:
	print("Fallback to file modified time.")
	# Use file mtime in local timezone and convert to UTC
	dt = datetime.fromtimestamp(os.path.getmtime(input_path), tz=local_tz).astimezone(timezone.utc)

# Injection
if not hdr_enabled:
	if profile_path:
		insert_before_first_idat(png, make_iccp_chunk(profile_path))
	insert_before_first_idat(png, make_chrm_chunk())
insert_before_first_idat(png, make_time_chunk(dt))
if process_name:
	insert_before_first_idat(png, make_text_chunk("Source", process_name))

# Overwrite input file
# Write output safely to temp then replace original
out_abs = os.path.abspath(input_path)
dirn = os.path.dirname(out_abs) or "."
tmp_path = os.path.join(dirn, os.path.basename(out_abs) + ".tmp")
with open(tmp_path, "wb") as f:
	f.write(png)
os.replace(tmp_path, out_abs)

print(f"Added metadata for {out_abs}!")