#!/usr/bin/env python3

from locale import getdefaultlocale
from struct import unpack
from collections import namedtuple
from typing import Any, Tuple

# https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-shllink/16cb4ca1-9339-4d0c-a68d-bf1d6cc0f943

TShellLink = namedtuple('ShellLink', ('ShellLinkHeader', 'LinkTargetIDList', 'LinkInfo', 'StringData', 'ExtraData'))
def parse_shell_link(data: bytes, offset: int=0) -> TShellLink:
	shell_link_header = parse_shell_link_header(data, offset)
	offset += shell_link_header.HeaderSize
	if shell_link_header.LinkFlags & 0x01:
		link_target_id_list = parse_link_target_id_list(data, offset)
		offset += link_target_id_list.IDListSize + 0x02
	else:
		link_target_id_list = None
	if shell_link_header.LinkFlags & 0x02:
		link_info = parse_link_info(data, offset)
		offset += link_info.LinkInfoSize
	else:
		link_info = None
	string_data = parse_string_data(data, offset, shell_link_header.LinkFlags)
	offset += string_data.StringDataSize
	return TShellLink(shell_link_header, link_target_id_list, link_info, string_data, parse_extra_data(data, offset))

TShellLinkHeader = namedtuple('ShellLinkHeader', ('HeaderSize', 'LinkCLSID', 'LinkFlags', 'FileAttributes', 'CreationTime', 'AccessTime', 'WriteTime', 'FileSize', 'IconIndex', 'ShowCommand', 'HotKey'))
def parse_shell_link_header(data: bytes, offset: int=0) -> TShellLinkHeader:
	sz = unpack('<I', data[offset: offset + 0x04])[0]
	return TShellLinkHeader(*unpack('<L16sLLQQQLlLH2x4x4x', data[offset: offset + sz]))

TLinkTargetIDList = namedtuple('LinkTargetIDList', ('IDListSize', 'IDList'))
def parse_link_target_id_list(data: bytes, offset: int=0) -> TLinkTargetIDList:
	sz = unpack('<H', data[offset: offset + 0x02])[0]
	return TLinkTargetIDList(*unpack('<H' + str(sz) + 's', data[offset: offset + 0x02 + sz]))

TLinkInfo = namedtuple('LinkInfo', ('LinkInfoSize', 'LinkInfoHeaderSize', 'LinkInfoFlags', 'VolumeIDOffset', 'LocalBasePathOffset', 'CommonNetworkRelativeLinkOffset', 'CommonPathSuffixOffset', 'LocalBasePathOffsetUnicode', 'CommonPathSuffixOffsetUnicode', 'VolumeID', 'LocalBasePath', 'CommonNetworkRelativeLink', 'CommonPathSuffix', 'LocalBasePathUnicode', 'CommonPathSuffixUnicode'))
TLinkInfo.__new__.__defaults__ = (None,)*len(TLinkInfo._fields)
def parse_link_info(data: bytes, offset: int=0) -> TLinkInfo:
	sz = unpack('<I', data[offset + 0x04: offset + 0x08])[0]
	res = TLinkInfo(*unpack('<' + 'L'*(sz//4), data[offset: offset + sz]))
	values = {}
	values['VolumeID'] = parse_blob(data, offset + res.VolumeIDOffset) if res.VolumeIDOffset > 0 else None
	values['LocalBasePath'] = parse_str(data, offset + res.LocalBasePathOffset) if res.LocalBasePathOffset > 0 else None
	values['CommonNetworkRelativeLink'] = parse_common_network_relative_link(data, offset + res.CommonNetworkRelativeLinkOffset) if res.CommonNetworkRelativeLinkOffset > 0 else None
	values['CommonPathSuffix'] = parse_str(data, offset + res.CommonPathSuffixOffset) if res.CommonPathSuffixOffset > 0 else None
	values['LocalBasePathUnicode'] = parse_ustr(data, offset + res.LocalBasePathOffsetUnicode) if res.LocalBasePathOffsetUnicode is not None and res.LocalBasePathOffsetUnicode > 0 else None
	values['CommonPathSuffixUnicode'] = parse_ustr(data, offset + res.CommonPathSuffixOffsetUnicode) if res.CommonPathSuffixOffsetUnicode is not None and res.CommonPathSuffixOffsetUnicode > 0 else None
	return res._replace(**values) if len(values) > 0 else res

TCommonNetworkRelativeLink = namedtuple('CommonNetworkRelativeLink', ('CommonNetworkRelativeLinkSize', 'CommonNetworkRelativeLinkFlags', 'NetNameOffset', 'DeviceNameOffset', 'NetworkProviderType', 'NetNameOffsetUnicode', 'DeviceNameOffsetUnicode', 'NetName', 'DeviceName', 'NetNameUnicode', 'DeviceNameUnicode'))
TCommonNetworkRelativeLink.__new__.__defaults__ = (None,)*len(TCommonNetworkRelativeLink._fields)
def parse_common_network_relative_link(data: bytes, offset: int=0) -> TCommonNetworkRelativeLink:
	sz = unpack('<I', data[offset + 0x08: offset + 0x0C])[0]
	res = TCommonNetworkRelativeLink(*unpack('<' + 'L'*(sz//4), data[offset: offset + sz]))
	values = {}
	values['NetName'] = parse_str(data, offset + res.NetNameOffset) if res.NetNameOffset > 0 else None
	values['DeviceName'] = parse_str(data, offset + res.DeviceNameOffset) if res.DeviceNameOffset > 0 else None
	values['NetNameUnicode'] = parse_ustr(data, offset + res.NetNameOffsetUnicode) if res.NetNameOffsetUnicode is not None and res.NetNameOffsetUnicode > 0 else None
	values['DeviceNameUnicode'] = parse_ustr(data, offset + res.DeviceNameOffsetUnicode) if res.DeviceNameOffsetUnicode is not None and res.DeviceNameOffsetUnicode > 0 else None
	return res._replace(**values) if len(values) > 0 else res

TStringData = namedtuple('StringData', ('StringDataSize', 'NAME_STRING', 'RELATIVE_PATH', 'WORKING_DIR', 'COMMAND_LINE_ARGUMENTS', 'ICON_LOCATION'))
TStringData.__new__.__defaults__ = (None,)*len(TStringData._fields)
def parse_string_data(data: bytes, offset: int=0, LinkFlags: int=0) -> TStringData:
	values = {}
	cursor = offset
	IsUnicode = LinkFlags & 0b10000000
	if LinkFlags & 0b100:
		num, values['NAME_STRING'] = parse_pstr(data, cursor, IsUnicode)
		cursor += num + 0x02
	if LinkFlags & 0b1000:
		num, values['RELATIVE_PATH'] = parse_pstr(data, cursor, IsUnicode)
		cursor += num + 0x02
	if LinkFlags & 0b10000:
		num, values['WORKING_DIR'] = parse_pstr(data, cursor, IsUnicode)
		cursor += num + 0x02
	if LinkFlags & 0b100000:
		num, values['COMMAND_LINE_ARGUMENTS'] = parse_pstr(data, cursor, IsUnicode)
		cursor += num + 0x02
	if LinkFlags & 0b1000000:
		num, values['ICON_LOCATION'] = parse_pstr(data, cursor, IsUnicode)
		cursor += num + 0x02
	return TStringData(StringDataSize=cursor - offset, **values)

def parse_extra_data(data: bytes, offset: int=0) -> Tuple[bytes, ...]:
	res = []
	while (sz := unpack('<I', data[offset: offset + 0x04])[0]) >= 0x04:
		res.append(data[offset: offset + sz])
		offset += sz
	return tuple(res)

def parse_str(data: bytes, offset: int=0) -> bytes:
	return data[offset:].split(b'\x00', 1)[0] if offset > 0 else data.split(b'\x00', 1)[0]

def parse_ustr(data: bytes, offset: int=0) -> bytes:
	return data[offset:].split(b'\x00\x00', 1)[0] if offset > 0 else data.split(b'\x00\x00', 1)[0]

def parse_pstr(data: bytes, offset: int=0, isunicode: bool=True) -> str:
	sz = unpack('<H', data[offset: offset + 0x02])[0]
	if isunicode:
		val = data[offset + 0x02: offset + 0x02 + 4*sz].decode('utf-16', errors='ignore')[:sz]
		return len(val.encode('utf-16-le')), val
	else:
		enc = getdefaultlocale()[1]
		val = data[offset + 0x02: offset + 0x02 + 4*sz].decode(enc, errors='ignore')[:sz]
		return len(val.encode(enc)), val

def parse_blob(data: bytes, offset: int=0) -> bytes:
	sz = unpack('<I', data[offset: offset + 0x04])[0]
	return data[offset: offset + sz]

def parse(x: Any) -> TShellLink:
	if isinstance(x, bytes):
		return parse_shell_link(x)
	elif hasattr(x, 'read') and callable(x.read):
		return parse_shell_link(x.read())
	elif isinstance(x, str):
		with open(x, 'rb') as stream:
			return parse_shell_link(stream.read())