from __future__ import print_function

"""
RCSid:
	$Id: tlv.py,v 1.8 2026/05/06 16:44:51 sjg Exp $

	@(#) Copyright (c) 2026 Simon J. Gerraty

	SPDX-License-Identifier: BSD-2-Clause
      
	Please send copies of changes and bug-fixes to:
	sjg@crufty.net

"""

import sys

# This works with python2, but the unit-test output looks odd
# we have to use bytearray in places where we would prefer
# bytes with python3
isPy2 = sys.version_info[0] == 2

def encode_int(n):
    """encode an integer as bytearray

    If value is < 0x80, a single byte is used.
    Otherwise the first byte indicates how many
    subsequent bytes (up to 4) hold the value.
    """
    if n < 0x80:
        ba = bytearray(1)
        ba[0] = n
    elif n < 0x10000:
        ba = bytearray(3)
        ba[0] = 0x82
        ba[1] = (n >> 8) & 0xff
        ba[2] = n & 0xff
    elif n < 0x1000000:
        ba = bytearray(4)
        ba[0] = 0x83
        ba[1] = (n >> 16) & 0xff
        ba[2] = (n >> 8) & 0xff
        ba[3] = n & 0xff
    else:
        ba = bytearray(5)
        ba[0] = 0x84
        ba[1] = (n >> 24) & 0xff
        ba[2] = (n >> 16) & 0xff
        ba[3] = (n >> 8) & 0xff
        ba[4] = n & 0xff

    if isPy2:
        return ba
    return bytes(ba)

def decode_int(ba):
    """decode a bytearray into an integer

    If the first byte has the 7th bit set,
    it indicates how many subsequent bytes hold the value
    So we consume 1,3,4 or 5 bytes

    return the value and the number of bytes consumed.
    """

    if isPy2:
        if isinstance(ba, str):
            ba = bytearray(ba[:5])
    n = int(ba[0])
    nbytes = n & 0x7f

    if n == nbytes:
        nbytes = 1
    else:
        nbytes += 1
        if len(ba) < nbytes:
            raise ValueError('Not enough bytes: expected {nbytes} bytes')
        if nbytes > 5:
            raise ValueError('Bad nbytes={nbytes}')
    
        if nbytes == 3:
            n = (ba[1] << 8) | ba[2]
        elif nbytes == 4:
            n = (ba[1] << 16) | (ba[2] << 8) | ba[3]
        else:
            n = (ba[1] << 24) | (ba[2] << 16) | (ba[3] << 8) | ba[4]
    return (n,nbytes)

class TLV(object):
    """Tag Length Value

    Tag integer
    Length integery
    Value bytes (or bytearray for python2)
    """

    encoding = "latin-1"

    def __init__(self, tag=None, value=None, length=-1, conf={}):
        self.conf = conf
        self.set_tag(tag)
        self.set_value(value)
        self.set_length(length)

    def set_tag(self, tag):
        if tag:
            self.tag = int(tag)
        else:
            self.tag = tag

    def set_value(self, value):
        """set value as bytes

        For python2 we have to use bytearray.
        """
        if not value:
            self.value = value
        else:
            if isinstance(value, bytes):
                self.value = value
            elif isPy2:
                self.value = bytearray(value)
            else:
                self.value = bytes(value, self.encoding)

    def set_length(self, length=-1):
        self.length = length
        if self.value:
            n = len(self.value)
            if length < 0:
                self.length = n

    def check_value(self):
        if self.value:
            n = len(self.value)
            if n != self.length:
                raise ValueError('Value length {} != {}: {}'.format(n, self.length, self))

    def encode(self):
        """return a sequence of bytes that represent this TLV

        For python2 we have to use bytearray.
        """
        btag = encode_int(self.tag)
        blen = encode_int(self.length)
        if isPy2:
            return bytearray(btag + blen + self.value)
        return btag + blen + self.value

    def decode(self, ba, conf={}):
        """decode TLV from ba and return bytes consumed

        If conf contains 'skip_value' we do not consume any bytes for
        the value - useful when caller only wants to know tag and
        length.

        Normally if we decode value will will run a sanity check that
        length of value matches that specified.  This is skipped if
        'skip_check' appears in conf.
        """
        (tag, xt) = decode_int(ba)
        (length, xl) = decode_int(ba[xt:xt+5])
        self.set_tag(tag)
        self.set_length(length)
        xs = xt + xl
        skip_value = conf.get('skip_value')
        if skip_value:
            return xs
        xe = xs + length
        # len(ba) could be expensive
        # we may get an IndexError here
        value = ba[xs:xe]
        # if we didn't it all, this may raise a ValueError
        self.set_value(value)
        skip_check = conf.get('skip_check')
        if not skip_check:
            self.check_value()
        return xe

    def __repr__(self):
        return "{},{},'{}'".format(self.tag,self.length,self.value)

    def __str__(self):
        if not self.value is None:
            sv = self.value.decode(self.encoding)
        else:
            sv = ''
        return "{},{},'{}'".format(self.tag,self.length,sv)

def encode(tlvs):
    """encode a list of TLVs to bytes"""
    ba = b''
    for t in tlvs:
        ba += t.encode()
    return ba

def decode(ba):
    """decode sequence of bytes into list of TLVs"""
    nbytes = len(ba)
    t = TLV()
    x = t.decode(ba)
    tlvs = [t]
    while x < nbytes:
        t = TLV()
        x += t.decode(ba[x:])
        tlvs.append(t)
    return tlvs

if __name__ == '__main__':
    for i in [64, 258, 0x10003, 0x1000004]:
        b = encode_int(i)
        (n,nb) = decode_int(b)
        print("i={} -> {} -> n={}".format(i,b,n))

    t1 = TLV(1, "Hello, World!")
    b1 = t1.encode()
    t2 = TLV()
    n = t2.decode(b1)
    print("{} -> {} -> {} {}".format(t1,b1,t2,n))

    t1 = TLV(1, "Hi")
    t2 = TLV(1, "there.")
    t3 = TLV(1, "We good?")

    ba = encode([t1,t2,t3])
    tlvs = decode(ba)
    for t in tlvs:
        print(t)

    te = TLV()
    n = len(b1)
    print("We expect an error")
    try:
        te.decode(b1[0:n-1])
    except IndexError as e:
        print(e)
    except ValueError as e:
        print(e)

    print("This time no error")
    te.decode(b1[0:n-1], {'skip_check': True})
    print(te)

    th = TLV()
    n = th.decode(b1, {'skip_value': True})
    print(th)
