"""
Script to calculate the Cyclomatic Complexity of every function.

Joxean Koret, 2019
Public Domain
"""

import idaapi

from tarjan_sort import strongly_connected_components

#-------------------------------------------------------------------------------
class CCFGDataChooser(idaapi.Choose2):
  def __init__(self, title, items):
    idaapi.Choose2.__init__(self,
                     title,
                     [ ["Address", 8], ["Name", 50], ["Nodes", 8],
                     ["Edges", 8], ["CC", 8], ["Loops", 8], ["Switchs", 8]])
    self.items = []
    for item in items:
      element = ["0x%08x" % item[0], item[1],
                 "%04d" % item[2], "%04d" % item[3], "%04d" % item[4],
                 "%04d" % item[5], "%03d" % item[6]]
      self.items.append(element)

  def OnSelectLine(self, n):
    ea = int(self.items[n][0], 16)
    jumpto(ea)

  def OnGetLine(self, n):
    return self.items[n]

  def OnGetSize(self):
    return len(self.items)

#-------------------------------------------------------------------------------
def analyze_function(ea):
  nodes = 0
  edges = 0
  cc = 0
  g = {}
  switches = set()

  func = idaapi.get_func(ea)
  if func is not None:
    nodes = 0
    edges = 0
    flow = idaapi.FlowChart(func)
    for block in flow:
      nodes += 1
      g[block.startEA] = []
      for succ in block.succs():
        edges += 1
        g[block.startEA].append(succ.startEA)

      for pred in block.preds():
        edges += 1
        try:
          g[pred.startEA].append(block.startEA)
        except KeyError:
          g[pred.startEA] = [block.startEA]

      if block.startEA == 0:
        continue

      for head in Heads(block.startEA, block.endEA):
        switch = idaapi.get_switch_info_ex(head)
        if switch:
          switches.add(switch.startea)

  cc = edges - nodes + 2
  return nodes, edges, cc, g, len(switches)

#-------------------------------------------------------------------------------
def main():
  vals = []
  items = []
  for ea in Functions():
    nodes, edges, cc, graph, switches = analyze_function(ea)
    try:
      scc = strongly_connected_components(graph)
      loops = 0
      for item in scc:
        val = len(item)
        if val > 1:
          loops += 1
    except:
      loops = -1

    name = GetFunctionName(ea)
    tmp = Demangle(name, INF_SHORT_DN)
    if tmp != "" and tmp is not None:
      name = tmp
    items.append([ea, name, nodes, edges, cc, loops, switches])
    vals.append(cc)

  line = "Cyclomatic complexity: Max %d, Avg %d"
  print(line % (max(vals), sum(vals)/len(vals)))
  
  ch = CCFGDataChooser("Control Flow Graphs Data", items)
  ch.Show()

if __name__ == "__main__":
  main()