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

Joxean Koret, 2019
Public Domain
"""

import idaapi

#-------------------------------------------------------------------------------
class CCyclomaticComplexityChooser(idaapi.Choose2):
  def __init__(self, title, items):
    idaapi.Choose2.__init__(self,
                     title,
                     [ ["Address", 8], ["Name", 50], ["Nodes", 8],
                     ["Edges", 8], ["CC", 8], ])
    self.items = []
    for item in items:
      element = ["0x%08x" % item[0], item[1],
                 "%04d" % item[2], "%04d" % item[3], "%04d" % item[4]]
      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 calc_cc(ea):
  func = idaapi.get_func(ea)
  if func is not None:
  
    nodes = 0
    edges = 0
    flow = idaapi.FlowChart(func)
    for block in flow:
      nodes += 1
      for succ in block.succs():
        edges += 1
      
      for pred in block.preds():
        edges += 1
  
  cc = edges - nodes + 2
  return nodes, edges, cc

#-------------------------------------------------------------------------------
def main():
  vals = []
  items = []
  for ea in Functions():
    nodes, edges, cc = calc_cc(ea)
    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])
    vals.append(cc)

  line = "Cyclomatic complexity: Max %d, Avg %d"
  print(line % (max(vals), sum(vals)/len(vals)))
  
  ch = CCyclomaticComplexityChooser("Cyclomatic Complexity", items)
  ch.Show()

if __name__ == "__main__":
  main()