from __future__ import print_function
import StringIO
import struct

class LineState:
  def __init__(self, is_stmt=False):
    self.address = 0
    self.file = 1
    self.line = 1
    self.column = 0
    self.is_stmt = is_stmt
    self.basic_block = False
    self.end_sequence = False
    self.section = None

  def clone(self):
    state = LineState()
    state.address = self.address
    state.file = self.file
    state.line = self.line
    state.column = self.column
    state.is_stmt = self.is_stmt
    state.basic_block = self.basic_block
    state.end_sequence = self.end_sequence
    state.section = self.section
    return state

  def __str__(self):
    return "[address=%08x file=%d line=%d column=%d is_stmt=%s basic_block=%s end_sequence=%d]" % \
      (self.address, self.file, self.line, self.column, self.is_stmt,
       self.basic_block, self.end_sequence)


class DwarfDebugLine:
  """decode .debug_line Dwarf line debug sections"""
  def __init__(self, verbose=False):
    self.input = None
    self.error = None
    self.verbose = verbose
    self.matrix = None

  def _log(self, *args):
    if self.verbose:
      print(*args)

  def decode(self, elf_file):
    # get section with debug info
    debug_line = elf_file.get_section_by_name(".debug_line")
    if debug_line is None:
      self.error = "No .debug_line section found! No debug info?"
      return False
    # get (optional) relocations
    rela = elf_file.get_section_by_name(".rela.debug_line")
    # start parsing
    self.input = StringIO.StringIO(debug_line.data)
    # decode header
    if not self.decode_header():
      return False
    if self.verbose:
      self.dump_header()
    # decode line program
    matrix = []
    state = LineState(self.default_is_stmt)
    log = self._log
    while True:
      # read opcode
      opc_ch = self.input.read(1)
      if len(opc_ch) == 0:
        break
      opc = ord(opc_ch)
      log("opcode=", opc)
      # 0 = extended opcode
      if opc == 0:
        opc_size = self.read_leb128()
        sub_opc = ord(self.input.read(1))
        log("  sub_opcode=", sub_opc)
        # 1: DW_LNE_end_sequence
        if sub_opc == 1:
          state.end_sequence = True
          line = state.clone()
          matrix.append(line)
          state.__init__()
          log("DW_LNE_end_sequence:", line)
        # 2: DW_LNE_set_address
        elif sub_opc == 2:
          pos = self.input.tell()
          addr = self.read_long()
          addend, sect = self.find_rela(rela, pos)
          state.address = addr + addend
          state.section = sect
          log("DW_LNE_set_address: %08x  sect=%s" % (state.address, sect))
        # 3: DW_LNE_set_file
        elif sub_opc == 3:
          tup = self.decode_file()
          self.files.append(tup)
          log("DW_LNE_set_file", tup)
        # other (unknown) ext opc
        else:
          log("unknown sub opcode!")
          self.input.seek(opc_size-1,1)
      # standard opcodes
      elif opc < self.opc_base:
        # 1: DW_LNS_copy
        if opc == 1:
          line = state.clone()
          matrix.append(line)
          log("DW_LNS_copy:", line)
          state.basic_block = False
        # 2: DW_LNS_advance_pc
        elif opc == 2:
          offset = self.read_leb128() * self.min_instr_len
          state.address += offset
          log("DW_LNS_advance_pc: +%d -> %08x" % (offset, state.address))
        # 3: DW_LNS_advance_line
        elif opc == 3:
          offset = self.read_sleb128()
          state.line += offset
          log("DW_LNS_advance_line: +%d -> %d" % (offset, state.line))
        # 4: DW_LNS_set_file
        elif opc == 4:
          state.file = self.read_leb128()
          log("DW_LNS_set_file", state.file)
        # 5: DW_LNS_set_column
        elif opc == 5:
          state.column = self.read_leb128()
          log("DW_LNS_set_column", state.column)
        # 6: DW_LNS_negate_stmt
        elif opc == 6:
          state.is_stmt = not state.is_stmt
          log("DW_LNS_negate_stmt", state.is_stmt)
        # 7: DW_LNS_set_basic_block
        elif opc == 7:
          state.basic_block = True
          log("DW_LNS_set_basic_block")
        # 8: DW_LNS_const_add_pc
        elif opc == 8:
          (addr_addend,_) = self.decode_special_opcode(255)
          state.address += addr_addend
          log("DW_LNS_const_add_pc: +%d -> %08x" % (addr_addend, state.address))
        # 9: DW_LNS_fixed_advance_pc
        elif opc == 9:
          offset = self.read_word()
          state.address += offset
          log("DW_LNS_fixed_advance_pc: %+08x" % offset)
        # other (unknown) opc
        else:
          num_args = self.std_opc_lens[opc]
          log("skip unknown: num_args=", num_args)
          for i in range(num_args):
            self.read_leb128()
      # special opcodes:
      else:
        (addr_addend, line_addend) = self.decode_special_opcode(opc)
        state.address += addr_addend
        state.line += line_addend
        state.basic_block = False
        line = state.clone()
        matrix.append(line)
        log("special", (opc - self.opc_base), line)
    # done
    self.matrix = matrix
    return True

  def get_matrix(self):
    return self.matrix

  def get_file_dir(self, idx):
    f = self.files[idx-1]
    dir_idx = f[1]
    if dir_idx > 0:
      dir_name = self.inc_dirs[dir_idx-1]
    else:
      dir_name = ""
    return dir_name

  def get_file_name(self, idx):
    return self.files[idx-1][0]

  def find_rela(self, rela_section, pos):
    if rela_section is not None:
      for rela in rela_section.rela:
        if rela.offset == pos:
          return rela.addend, rela.section
    return 0,None

  def decode_special_opcode(self, opc):
    adj_opc = opc - self.opc_base
    addr_addend = (adj_opc / self.line_range) * self.min_instr_len
    line_addend = self.line_base + (adj_opc % self.line_range)
    return (addr_addend, line_addend)

  def decode_header(self):
    # header
    self.unit_length = self.read_long()
    self.version = self.read_word()
    if self.version != 2:
      self.error = "Can only decode DWARF 2 debug info"
      return False
    self.header_length = self.read_long()
    self.min_instr_len = self.read_byte()
    self.default_is_stmt = self.read_byte()
    self.line_base = self.read_sbyte()
    self.line_range = self.read_byte()
    self.opc_base = self.read_byte()
    # 9 standard opcode lengths
    self.std_opc_lens = []
    n = self.opc_base
    if n > 0:
      for i in range(n-1):
        l = self.read_byte()
        self.std_opc_lens.append(l)
    # 10 include dirs
    self.inc_dirs = []
    while True:
      inc_dir = self.read_string()
      if inc_dir == "":
        break
      self.inc_dirs.append(inc_dir)
    # 11 file names
    self.files = []
    while True:
      tup = self.decode_file()
      if tup is None:
        break
      self.files.append(tup)
    # end header: check header size
    pos = self.input.tell()
    hdr_len = pos - 10
    if hdr_len != self.header_length:
      self.error = "Error size mismatch: %d != %d" % (hdr_len, self.header_length)
      return False
    return True

  def decode_file(self):
    file_name = self.read_string()
    if file_name == "":
      return None
    dir_idx = self.read_leb128()
    last_mod = self.read_leb128()
    file_size = self.read_leb128()
    return (file_name, dir_idx, last_mod, file_size)

  def dump_header(self):
    print("unit_length=%x version=%d header_length=%x max_instr_len=%d "
          "default_is_stmt=%d line_base=%d line_range=%d opc_base=%d" %
          (self.unit_length, self.version, self.header_length,
           self.min_instr_len, self.default_is_stmt, self.line_base,
           self.line_range, self.opc_base))
    print("std_opc_lens:",",".join(map(str,self.std_opc_lens)))
    print("inc_dirs")
    for d in self.inc_dirs:
      print(d)
    print("files")
    for f in self.files:
      print(f)

  def read_string(self):
    result = []
    while True:
      ch = self.input.read(1)
      if ord(ch) == 0:
        break
      result.append(ch)
    return "".join(result)

  def read_leb128(self):
    result = 0
    shift = 0
    while True:
      byte = self.read_byte()
      result |= (byte & 0x7f) << shift
      if byte & 0x80 == 0:
        break
      shift += 7
    return result

  def read_sleb128(self):
    result = 0
    shift = 0
    while True:
      byte = self.read_byte()
      result |= (byte & 0x7f) << shift
      shift += 7
      if byte & 0x80 == 0:
        break
    # negative?
    if byte & 0x40 == 0x40:
      mask = 1 << shift
      result |= - mask
    return result

  def read_long(self):
    data = self.input.read(4)
    return struct.unpack(">I", data)[0]

  def read_word(self):
    data = self.input.read(2)
    return struct.unpack(">H", data)[0]

  def read_byte(self):
    data = self.input.read(1)
    return struct.unpack(">B", data)[0]

  def read_sbyte(self):
    data = self.input.read(1)
    return struct.unpack(">b", data)[0]

# mini test
if __name__ == '__main__':
  import sys
  from ELFReader import ELFReader
  reader = ELFReader()
  for a in sys.argv[1:]:
    f = open(a, "rb")
    ef = reader.load(f)
    ddl = DwarfDebugLine(verbose=True)
    ok = ddl.decode(ef)
    if ok:
      print("--- line matrix ---")
      for row in ddl.get_matrix():
        name = ddl.get_file_name(row.file)
        fdir = ddl.get_file_dir(row.file)
        sect_name = row.section.name_str
        print("%08x: %s [%s] %s:%d" % (row.address, sect_name, fdir, name, row.line))
