from laser.ethereum import svm
from z3 import Z3Exception, simplify
import re


graph_html = '''<html>
 <head>
  <style type="text/css">
   #mynetwork {
    background-color: #232625;
   }

   body {
    background-color: #232625;
    color: #ffffff;
   }
  </style>
  <link href="https://cdnjs.cloudflare.com/ajax/libs/vis/4.21.0/vis.min.css" rel="stylesheet" type="text/css" />
  <script src="https://cdnjs.cloudflare.com/ajax/libs/vis/4.21.0/vis.min.js"></script>
  <script>

    var options = {
      autoResize: true,
      height: '100%',
      width: '100%',
      manipulation: false,
      height: '90%',
      layout: {
        randomSeed: undefined,
        improvedLayout:true,
        hierarchical: {
          enabled:true,
          levelSeparation: 450,
          nodeSpacing: 200,
          treeSpacing: 100,
          blockShifting: true,
          edgeMinimization: true,
          parentCentralization: false,
          direction: 'LR',        // UD, DU, LR, RL
          sortMethod: 'directed'   // hubsize, directed
        }
      },
      nodes:{
        borderWidth: 1,
        borderWidthSelected: 2,
        chosen: true,
        shape: 'box',
        font: {
          align: 'left',
          color: '#FFFFFF',
        },
      },
      edges:{
        font: {
          color: '#ffffff',
          size: 12, // px
          face: 'arial',
          background: 'none',
          strokeWidth: 0, // px
          strokeColor: '#ffffff',
          align: 'horizontal',
          multi: false,
          vadjust: 0,
        }
      },

      physics:{
        enabled: [ENABLE_PHYSICS],
      }   
    
  }

  [JS]

  </script>
 </head>
<body>
<p>Mythril / Ethereum LASER Symbolic VM</p>
<p><div id="mynetwork"></div><br /></p>
<script type="text/javascript">
var container = document.getElementById('mynetwork');
var data = {'nodes': nodes, 'edges': edges}
var gph = new vis.Network(container, data, options);
</script>
</body>
</html>
'''


colors = [
  "{border: '#26996f', background: '#2f7e5b', highlight: {border: '#26996f', background: '#28a16f'}}",
  "{border: '#9e42b3', background: '#842899', highlight: {border: '#9e42b3', background: '#933da6'}}",
  "{border: '#b82323', background: '#991d1d', highlight: {border: '#b82323', background: '#a61f1f'}}",
  "{border: '#4753bf', background: '#3b46a1', highlight: {border: '#4753bf', background: '#424db3'}}",
]  


def serialize(_svm, color_map):

    nodes = []
    edges = []

    print(_svm.nodes)

    for n in _svm.nodes:

        code =  _svm.nodes[n].as_dict()['code']

        code = re.sub("([0-9a-f]{8})[0-9a-f]+",  lambda m: m.group(1) + "(...)", code)

        color = color_map[_svm.nodes[n].as_dict()['module_name']]

        nodes.append("{id: '" + str(_svm.nodes[n].as_dict()['id']) + "', color: " + color + ", size: 150, 'label': '" + code + "'}")

    for edge in _svm.edges:

      if (edge.condition is None):
          label = ""
      else:

          try:
            label = str(simplify(edge.condition)).replace("\n", "")
          except Z3Exception:
            label = str(edge.condition).replace("\n", "")
    
      label = re.sub("([^_])([\d]{2}\d+)",  lambda m: m.group(1) + hex(int(m.group(2))), label)
      code = re.sub("([0-9a-f]{8})[0-9a-f]+",  lambda m: m.group(1) + "(...)", code)

      edges.append("{from: '" + str(edge.as_dict()['from']) + "', to: '" + str(edge.as_dict()['to']) + "', 'arrows': 'to', 'label': '" + label + "', 'smooth': {'type': 'cubicBezier'}}")

    return "var nodes = [\n" + ",\n".join(nodes) + "\n];\nvar edges = [\n" + ",\n".join(edges) + "\n];"



def generate_callgraph(modules, physics):

    _svm = svm.SVM(modules)

    _svm.sym_exec()

    i = 0

    color_map = {}

    for k in modules:
      color_map[modules[k]['name']] = colors[i]
      i += 1

    html = graph_html.replace("[JS]", serialize(_svm, color_map))
    html = html.replace("[ENABLE_PHYSICS]", str(physics).lower())

    return html

@@ -2463,36 +2463,20 @@
     
-print(_svm.nodes)%0A%0A    for n
+for node_key
  in 
@@ -2508,32 +2508,39 @@
  =  _svm.nodes%5Bn
+ode_key
 %5D.as_dict()%5B'cod
@@ -2662,32 +2662,39 @@
 map%5B_svm.nodes%5Bn
+ode_key
 %5D.as_dict()%5B'mod
@@ -2742,42 +2742,16 @@
 %22 + 
-str(_svm.nodes%5Bn%5D.as_dict()%5B'id'%5D)
+node_key
  + %22

Update CFG generator to new node format