from scapy import *

# New objects to support TLV

class GenericCountField(Field): # TODO: make this class lighter
  def i2m(self, pkt, x):
    if x is None:
      x = getattr(pkt,self.field)
      i = 0
      while isinstance(x, self.cls): 
	x = x.payload
	i += 1
      x = i
    return x
  def i2h(self, pkt, x):
    return self.i2m(pkt, x)

class ByteCountField(GenericCountField, ByteField): 
  def __init__(self, name, default, field, cls):
    ByteField.__init__(self, name, default)
    self.field = field
    self.cls = cls

class ShortCountField(GenericCountField, ShortField):
  def __init__(self, name, default, field, cls):
    ShortField.__init__(self, name, default)
    self.field = field
    self.cls = cls

class MainCountField(StrField):
  def __init__(self, name, countfld, passon=0):
    StrField.__init__(self, name, None)
    self.countfld = countfld
    self.passon = passon

  def i2m(self, pkt, x):
    if x is None:
	return ""
    return str(x)
  
  def dissectme(self, s, p):
    return None,p # TODO: ?!?

  def getfield(self, pkt, s):
    if type(s) is tuple :
      s,p = s
    else:
      p = 0
    ret = None
    c = getattr(pkt, self.countfld)
    while c:
      c -= 1
      r,p = self.dissectme(s, p)
      if ret is None:
	ret = r
      else:
	ret.add_payload(r)
    if self.passon:
      return (s,p),ret
    else:
      return s[p:],ret

# Here comes generic code

class TLVPacket(Packet):
  name = 'TLV'
  fields_desc = [ ByteField('type', 0),
                  ByteField('len', None),
                  StrField('value', None)
                ]

class TLVCountField(ShortCountField):
  def __init__(self, name, default, field):
    ShortCountField.__init__(self, name, default, field, TLVPacket)

class TLVField(MainCountField):
  def __init__(self, name, countfld):
    MainCountField.__init__(self, name, countfld)
  # Note: each instance of MainCountField should re-implement dissectme()
  def dissectme(self, s, p):
    t = s[p:p+2]
    type,l = struct.unpack('BB', t)
    tlv = TLVPacket(s[p:p+2+l])
    p = p+l+2
    return tlv,p

class Header(Packet):
  name = 'TLV based header'
  fields_desc = [ 
                  TLVCountField('num', None, 'tlv'),
                  TLVField('tlv', 'num')
                 ]

if __name__ == '__main__':
  tlv=Header(num=1, tlv=TLVPacket(type=1, len=4, value='toto'))
  h=Header()
  h.dissect(str(tlv))
  h.display()

  tlv=Header(num=2, tlv=TLVPacket(type=1, len=4, value='toto')/TLVPacket(type=2, len=5, value='pouet'))
  h=Header()
  h.dissect(str(tlv))
  h.display()