def IsRegularGraph(topo: nx.DiGraph) -> bool:
"""Dtermine a graph is regular or not, i.e. all nodes have the same degree."""
degree = topo.degree(0)
for rank in range(1, topo.number_of_nodes()):
if topo.degree(rank) != degree:
return False
return True
def compute_centrality(graph: nx.DiGraph,
data: geojson.feature.FeatureCollection,
edge_map: Dict):
for item in data['features']:
edge = edge_map[item['properties']['id']]
from_degree = graph.degree(edge[0])
to_degree = graph.degree(edge[1])
item['properties']["from_degree"] = from_degree
item['properties']["to_degree"] = to_degree
def remove_unconnected_nodes(graph: nx.DiGraph):
"""Remove unconnected nodes and return the corresponding issues."""
removed_nodes = [node for node in graph.nodes() if graph.degree(node) == 0]
removed_issues = {graph.node[node]['issue'] for node in removed_nodes}
if removed_nodes:
graph.remove_nodes_from(removed_nodes)
return removed_issues
def create_node_data(node: NodeTy, commit: CommitRepoPair,
cig: nx.DiGraph) -> ArcPlotNodeInfo:
return {
"info": commit.commit_hash.short_hash,
"size": cig.degree(node),
"fill_color": cig.out_degree(node),
"line_color": cig.in_degree(node)
}
def calc_node_size(graph: nx.DiGraph, node: tg.Text) -> int:
"""Size nodes by their degree.
Use natrual log to make them less diversed
multiply by SIZE_FACTOR to make them not too small
plus 1 so that size won't be zero
"""
SIZE_FACTOR = 3
return math.floor(SIZE_FACTOR * math.log1p(graph.degree(node))) + 1
def _construct(self):
"""
We want to build the type of DFG that's used in "Automated Ident. of Crypto
Primitives in Binary Code with Data Flow Graph Isomorphisms." Unlike that
paper, however, we're building it on Vex IR instead of assembly instructions.
"""
cfg = self._cfg
p = self.project
dfgs = {}
l.debug("Building Vex DFG...")
for node in cfg.nodes():
try:
if node.simprocedure_name == None:
irsb = p.factory.block(node.addr).vex
else:
l.debug("Cannot process SimProcedures, ignoring %s" %
node.simprocedure_name)
continue
except Exception as e:
l.debug(e)
continue
tmpsnodes = {}
storesnodes = {}
putsnodes = {}
statements = irsb.statements
dfg = DiGraph()
for stmt_idx, stmt in enumerate(statements):
# We want to skip over certain types, such as Imarks
if self._need_to_ignore(node.addr, stmt, stmt_idx):
continue
# break statement down into sub-expressions
exprs = stmt.expressions
stmt_node = stmt
dfg.add_node(stmt)
if stmt.tag == 'Ist_WrTmp':
tmpsnodes[stmt.tmp] = stmt_node
if exprs[0].tag == 'Iex_Binop':
if exprs[1].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[1].tmp], stmt_node)
else:
dfg.add_edge(exprs[1], stmt_node)
if exprs[2].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[2].tmp], stmt_node)
else:
dfg.add_edge(exprs[2], stmt_node)
elif exprs[0].tag == 'Iex_Unop':
dfg.remove_node(stmt_node)
if exprs[1].tag == 'Iex_RdTmp':
tmpsnodes[stmt.tmp] = copy(tmpsnodes[exprs[1].tmp])
tmpsnodes[stmt.tmp].tmp = stmt.tmp
else:
tmpsnodes[stmt.tmp] = exprs[1]
elif exprs[0].tag == 'Iex_RdTmp':
tmpsnodes[stmt.tmp] = copy(tmpsnodes[exprs[0].tmp])
tmpsnodes[stmt.tmp].tmp = stmt.tmp
elif exprs[0].tag == 'Iex_Get':
if putsnodes.has_key(exprs[0].offset):
dfg.add_edge(putsnodes[exprs[0].offset], stmt_node)
if len(exprs) > 1 and exprs[1].tag == "Iex_RdTmp":
dfg.add_edge(tmpsnodes[exprs[1].tmp], stmt_node)
elif len(exprs) > 1:
dfg.add_edge(exprs[1], stmt_node)
elif exprs[0].tag == 'Iex_Load':
if exprs[1].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[1].tmp], stmt_node)
else:
dfg.add_edge(exprs[1], stmt_node)
else:
# Take a guess by assuming exprs[0] is the op and any other expressions are args
for e in exprs[1:]:
if e.tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[e.tmp], stmt_node)
else:
dfg.add_edge(e, stmt_node)
elif stmt.tag == 'Ist_Store':
if exprs[0].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[0].tmp], stmt_node)
elif exprs[0].tag == 'Iex_Const':
dfg.add_edge(exprs[0], stmt_node)
if exprs[1].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[1].tmp], stmt_node)
else:
dfg.add_edge(exprs[1], stmt_node)
elif stmt.tag == 'Ist_Put':
if exprs[0].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[0].tmp], stmt_node)
elif exprs[0].tag == 'Iex_Const':
dfg.add_edge(exprs[0], stmt_node)
putsnodes[stmt.offset] = stmt_node
elif stmt.tag == 'Ist_Exit':
if exprs[0].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[0].tmp], stmt_node)
elif stmt.tag == 'Ist_Dirty':
tmpsnodes[stmt.tmp] = stmt_node
elif stmt.tag == 'Ist_CAS':
tmpsnodes[stmt.oldLo] = stmt_node
else:
for e in stmt.expressions:
if e.tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[e.tmp], stmt_node)
else:
dfg.add_edge(e, stmt_node)
for vtx in dfg.nodes():
if dfg.degree(vtx) == 0:
dfg.remove_node(vtx)
if dfg.size() > 0:
dfgs[node.addr] = dfg
return dfgs
class DepF :
def __init__(self, field) :
self.__offsets = []
self.__field = field
############ create depency field graph #########
# Initial values to build the graph (depth, width, cycles)
self.__depth = 3
self.__width = 3
self.__cycles = 2
self.__G = DiGraph()
G = self.__G
G.add_node( self._new_node(G) )
# Create randomlu the graph without cycle
self.__random( G, 0 )
# Find the best path to add cycles
d = all_pairs_dijkstra_path_length( G )
l = list( reversed( sorted( d, key=lambda key: len(d[key]) ) ) )
for i in l :
if self.__cycles == 0 :
break
d_tmp = sorted( d[i], key=lambda key: d[i][key] )
G.add_edge( d_tmp[-1], i)
self.__cycles = self.__cycles - 1
print simple_cycles( G )
print G.node
print G.edge
print G.degree()
############ Add field <-> higher #####################################
# field F depends of the higher degree
# F <---> sorted( G.degree(), key=lambda key : G.degree()[key] )[-1]
#######################################################################
degree = G.degree()
high_degree = sorted( degree, key=lambda key : degree[key] )[-1]
# Link our protected field with the node which has the highest degree
G.add_edge( field.get_name(), high_degree )
G.add_edge( high_degree, field.get_name() )
#draw_graphviz(G)
#write_dot(G,'file.dot')
def get_field(self) :
return self.__field
def add_offset(self, idx) :
x = Offset( idx )
self.__offsets.append( x )
return x
def run(self, _vm, _analysis, _vm_generate) :
###############################################################
## dict (method) of list ( offset / list (of instructions) ) #
## - insert an element #
## - modify the offset with the new insertion #
###############################################################
list_OB = {}
############ Create dependencies fields ############
fields = { self.__field.get_name() : Field( _vm, _analysis, _vm_generate, self.__field, self, True ) }
fields[ self.__field.get_name() ].run( self.__G.degree()[ self.__field.get_name() ] )
print self.__field.get_name(), ": REAL_FIELD ->", self.__G.predecessors( self.__field.get_name() ), self.__G.degree()[self.__field.get_name()]
############ Create the name, the initial value and all access of the field ############
for i in self.__G.node :
# We have not yet add this new field
if i not in fields :
print i, "PRE ->", self.__G.predecessors( i ), self.__G.degree()[i]
name, access_flag, descriptor = _analysis.get_like_field()
_vm.insert_field( self.__field.get_class_name(), name, [ access_flag, descriptor ] )
fields[ i ] = Field( _vm, _analysis, _vm_generate, _vm.get_field_descriptor( self.__field.get_class_name(), name, descriptor ), self )
# degree of the field, prefix must be add if the protection is for a real field
fields[ i ].run( self.__G.degree()[i] )
########## Add all fields initialisation into the final list ############
for i in fields :
print "FIELD ->", i,
fields[ i ].show()
x = fields[ i ].insert_init()
if x != None :
try :
list_OB[ x[0] ].append( (x[1], x[2]) )
except KeyError :
list_OB[ x[0] ] = []
list_OB[ x[0] ].append( (x[1], x[2]) )
############ Create the depedency ############
############################################################################
# Integer variables :
# X -> Y
# - a depth into the calcul
# Y = { LV + LLV + NLV } x { &, -, +, |, *, /, ^ } x { &&, || }
# if (Y != ?) {
# X = ..... ;
# }
#############################################################################
# get a local variable
# - used into a loop
# - a parameter
# - a new one
############################################################################
find = False
print "F -->", self.__G.successors( self.__field.get_name() )
taint_field = _analysis.get_tainted_field( self.__field.get_class_name(), self.__field.get_name(), self.__field.get_descriptor() )
for path in taint_field.get_paths() :
continue
print "\t", path.get_access_flag(), "%s (%d-%d)" % (path.get_bb().get_name(), path.get_bb().get_start(), path.get_bb().get_end()) , path.get_bb().get_start() + path.get_idx(),
x = _analysis.get( path.get_bb().get_method() )
bb = x.get_break_block( path.get_bb().get_start() + path.get_idx() )
print "\t\t", x.get_local_variables()
if path.get_access_flag() == "R" and find == False :
o = self.add_offset( _analysis.prev_free_block_offset( path.get_method(), bb.get_start() ) )
val = _analysis.next_free_block_offset( path.get_method(), o.get_idx() )
if o.get_idx() == -1 or val == -1 :
raise("ooop")
#try :
# list_OB[ path.get_method() ].append( o, [ [ "iload_3" ], [ "iconst_0" ], [ "if_icmpge", val - o.get_idx() + 3 ] ] )
#except KeyError :
# list_OB[ path.get_method() ] = []
# list_OB[ path.get_method() ].append( (o, [ [ "iload_3" ], [ "iconst_0" ], [ "if_icmpge", val - o.get_idx() + 3 ] ] ) )
find = True
##### Insert all modifications
for m in list_OB :
code = m.get_code()
i = 0
while i < len( list_OB[ m ] ) :
v = list_OB[ m ][ i ]
print "INSERT ", v[0].get_idx(), v[1]
size_r = code.inserts_at( code.get_relative_idx( v[0].get_idx() ), v[1] )
#code.show()
j = i + 1
while j < len( list_OB[ m ] ) :
v1 = list_OB[ m ][ j ]
if v1[0].get_idx() >= v[0].get_idx() :
v1[0].add_idx( size_r )
j = j + 1
i = i + 1
def _new_node(self, G) :
return "X%d" % (len(G.node))
def _current_node(self, G) :
return len(G.node) - 1
def __random(self, G, depth) :
if depth >= self.__depth :
return
for i in range( random.randint(1, self.__width) ) :
nd = self._new_node(G)
G.add_edge( "X%d" % depth, nd )
self.__random( G, self._current_node(G) )
class ProcessGraph(object):
"""
Helper class to represent a materials synthesis action graph / quest in a
format facilitating easy visualization in process form. This is also used
by the SurfaceGenerator to create the textual description of a quest.
"""
def __init__(self) -> None:
self.G = DiGraph()
# For each material we hold a list of nodes representing it over time
# as its state changes
self.ent_states_map = {}
def init_ent(self, e: Variable):
""" Create and initialize new EntityNode. """
if not e in self.ent_states_map:
self.ent_states_map[e] = [EntityNode(var=e, g=0)]
def curr_state(self, e: Variable):
""" Return the last (current) state for a given entity variable. """
return self.ent_states_map[e][-1]
def add_new_state(self, e: Variable):
"""
Add new state for a given entity variable (except Operation entities).
This creates a new EntityNode with g increased by 1.
"""
if KnowledgeBase.default().types.is_descendant_of(e.type, ['op']):
return self.ent_states_map[e]
else:
new_state_node = EntityNode(var=e, g=(self.curr_state(e).g + 1))
self.ent_states_map[e].append(new_state_node)
return new_state_node
def topological_sort_actions(self):
"""
Return edges in "topological order", meaning we iterate over nodes
in topological order and for each node iterate over its outgoing
edges.
"""
edge_list = []
for n in nx.topological_sort(self.G):
edge_list += edges_to_actions(self.G.out_edges(n, data=True))
return edge_list
def get_incoming_actions(self, node: EntityNode) -> List[ActionEdge]:
""" Return the incoming actions to a given EntityNode. """
return edges_to_actions(self.G.in_edges(node, data=True))
def get_actions_path(self, source: EntityNode,
target: EntityNode) -> List[ActionEdge]:
"""
Return list of ActionEdges comprising the shortest path between a
source and target node in the graph.
"""
action_path = []
path_nodes = nx.shortest_path(self.G, source, target)
if len(path_nodes) < 2:
return action_path
edges = [e for e in zip(path_nodes[:-1], path_nodes[1:])
] # get edges comprising path
for u, v in edges:
action_path.append(
ActionEdge(action=self.G.edges()[(u, v)]['action'],
source=u,
target=v))
return action_path
def get_start_material_nodes(self) -> List[EntityNode]:
""" Get nodes of all materials at start of process (g=0). """
return [
n for n in self.G.nodes() if (n.var.type == 'm') and (n.g == 0)
]
def get_start_material_vars(self) -> List[Variable]:
""" Get variables of all materials at start of process (g=0). """
return [mat.var for mat in self.get_start_material_nodes()]
def from_tw_actions(self, actions: Iterable[Action]) -> DiGraph:
""" Construct internal graph given a TextWorld Quest. """
# TODO verify validity?
from tw_textlabs.generator.quest_generator import convert_to_compact_action
G = nx.DiGraph()
action_vars = [convert_to_compact_action(a) for a in actions]
for action_var in action_vars:
action_name = action_var.name
if action_name == 'op_o_obtain':
# add change state var as arg to this action so it gets processed
action_var.vars.insert(0, action_var.change_state_vars[0])
ents = action_var.vars
for e in ents:
self.init_ent(e)
# ignore actions with single arg
if len(ents) != 2:
continue
is_state_change_action = len(action_var.change_state_vars) > 0
if not is_state_change_action:
G.add_edge(self.curr_state(ents[0]),
self.curr_state(ents[1]),
action=action_name)
# assuming just one material changes state per action
else:
assert (len(action_var.change_state_vars) == 1)
v = action_var.change_state_vars[0]
assert (v == ents[0])
# second argument is the source for new one (such as device)
source = self.curr_state(ents[1])
# Add a new state for the device, we assume that each operation
# on a material is at a new state.
if self.curr_state(v) in G.nodes():
# entity already exists in graph
target = self.add_new_state(v)
else:
# entity doesn't exist yet in graph, connect init node
target = self.curr_state(v)
# Add 'result' type edge which corresponds to 'obtain' action
G.add_edge(source, target, action='obtain')
# add Generator dummy node to be source of graph (cosmetic)
self.generator = EntityNode(var=Variable(name='START',
type=GENERATOR_DUMMY_TYPE),
g=0)
G.add_node(self.generator)
self.G = G
materials = self.get_start_material_nodes()
for m in materials:
self.G.add_edge(self.generator, m, action='take')
def get_op_type(self, node: EntityNode) -> str:
assert (KnowledgeBase.default().types.is_descendant_of(
node.var.type, 'tlq_op'))
op_type_node = [
n for n in self.G.neighbors(node)
if KnowledgeBase.default().types.is_descendant_of(
n.var.type, 'toe')
]
assert (len(op_type_node) == 1)
op_type_node = op_type_node[0]
return op_type_node.var.type
def get_source_op(self, node: EntityNode) -> EntityNode:
assert (KnowledgeBase.default().types.is_descendant_of(
node.var.type, ['m', 'sa']))
# return first operation node producing target node
for n in self.ent_states_map[node.var]:
for s, t in self.G.in_edges(n):
if s.var.type == 'tlq_op':
return s
return None
def rename_edges(self,
node: EntityNode,
orig: str,
target: str,
incoming: bool = True) -> None:
""" Rename a node's incoming or outgoing edges """
edges = [e for e in self.G.in_edges(node)
] if incoming else [e for e in self.G.out_edges(node)]
for edge in edges:
act = self.G.edges()[edge]['action']
self.G.edges()[edge]['action'] = act.replace(orig, target)
def draw(self):
""" Draw the ProcessGraph. """
connected_nodes = [n for n in self.G.nodes() if self.G.degree(n) > 0]
connected_G = self.G.subgraph(connected_nodes)
# Hierarchical layout
pos = graphviz_layout(connected_G, prog='dot')
edge_labels = dict([((
u,
v,
), d['action']) for u, v, d in self.G.edges(data=True)])
node_labels = {n: str(n) for n in connected_G.nodes()}
# Add dummy source node
node_labels.update({self.generator: 'START'})
node_colors = [node_to_color(n) for n in connected_G.nodes()]
nx.draw(connected_G,
pos,
labels=node_labels,
with_labels=True,
node_color=node_colors)
nx.draw_networkx_edge_labels(connected_G, pos, edge_labels=edge_labels)
class Case3(Case):
"""
This class defines a radial network which will be studied further. It uses directed graph implemented in networkx.
"""
def __init__(self,
name=None,
base_mva=1.0,
base_kv=20.0,
buses=None,
branches=None,
generators=None):
#Case.__init__(self,buses=buses,branches=branches)
self.buses = buses
self.base_mva = base_mva
self.base_kv = base_kv
self.base_z = base_kv**2 / base_mva
self.branches = branches
self.generators = generators
self.G = DiGraph()
self._buildGraph()
self.rootbus = [b for b in self.buses if self.G.pred[b] == {}]
self.forkbuses = [b for b in self.buses if self.G.degree()[b] > 2]
self.normalbuses = [b for b in self.buses if self.G.degree()[b] == 2]
self.endbuses = [b for b in self.buses if self.G.succ[b] == {}]
self.updateTypeG()
self.bus_level = shortest_path_length(self.G, source=self.rootbus[0])
self.updateBusConnBranch()
self.buses_loadJ = self.updateBusLoadCurrent()
self.identifyPVbuses()
self.calculateZpos()
for b in self.buses:
b.base_mva = base_mva
#draw(self.G)
#pl.show()
def _buildGraph(self):
""" add buses as nodes and branches as edges.
"""
self.G.add_nodes_from([bus for bus in self.buses])
self.G.add_edges_from([(b.from_bus, b.to_bus) for b in self.branches])
def plotGraph(self):
draw(self.G)
show()
return
def plotVP(self):
""" Plot voltage profile of distribution system"""
V = [abs(b.E) for b in self.buses]
name = [b.name for b in self.buses]
plot(V)
show()
return
def updateBusConnBranch(self):
"""
Updating to which branch a bus is connected.
"""
for a in self.branches:
a.from_bus.connected_to_branch.append(a)
a.to_bus.connected_from_branch.append(a)
# as edges are added, there should be a procedure to update bus.connected_to and bus.connected_from
def updateTypeG(self):
for b in self.buses:
if self.rootbus.__contains__(b):
b.typeG = rootNode
if self.forkbuses.__contains__(b):
b.typeG = forkNode
if self.endbuses.__contains__(b):
b.typeG = endNode
def updateBusLoadCurrent(self):
return [b.calc_loadI() for b in self.buses]
def getRoot(self):
"""
Get the root node of a network which is the substation node.
"""
rootNode = [n for n in self.G.node if self.G.pred[n] == {}]
return rootNode[0]
def updateBaseKV(self, aNode):
"""
update base_kv of aNode children. This is done after a voltage regulator changes its tap.
"""
for b in aNode.connected_to_branch:
b.updateToBusBaseKV()
if b.to_bus.connected_to_branch != []:
self.updateBaseKV(b.to_bus)
def identifyPVbuses(self):
"""
identify pv buses in the network. Important when creating the positive sequence matrix needed to update injected current
"""
self.pvbuses = [
g.bus for g in self.generators if g.type == 'PV'
] # a list of buses connected to generators controlled as PV bus.
for b in self.pvbuses:
b.type = 'PV' # update the bus.type according to its generator control mode.
def calculateZpos(self):
"""
calculate positive impedance matrix for updating injected current at PV bus in BFS with DGs
"""
self.identifyPVbuses()
pathToRoot = {}
lines = {}
Z1diag = {}
pvbusindex = {}
d = defaultdict(dict)
npvb = len(self.pvbuses) # number of pv buses
Z1 = zeros([npvb, npvb], dtype='complex')
for pvb in self.pvbuses:
pathToRoot[pvb] = shortest_path(self.G, self.rootbus[0], pvb)
lines[pvb] = [
b.connected_from_branch[0] for b in pathToRoot[pvb]
if b.connected_from_branch != []
]
self.pathToRoot = pathToRoot # a list of buses that incident with the path from a pv bus to root bus.
self.lines = lines # a list of lines that forms the path from a pvbus to the root bus.
Z_1 = zeros([npvb,
npvb]) # prepare an all zero matrix with size npvb x npvb
# This function requires a positif sequence of all lines in self.lines. The positive sequence impedance is available at line.Z012[0]
# the situation is more intricate if there is a transformer within the path between the PV bus to root bus. If on of the transformer winding
# is connected in delta, the path for positive sequence is disconnected. The implementation for this kind of situation is postponed.
# the current implementation assumes that there is NO TRANSFORMER exist in the path. 27th May 2011.
for pvb in self.pvbuses:
Z1sum = complex(0, 0)
for l in self.lines[pvb]:
Z1sum = Z1sum + l.Z012[1][1]
Z1diag[pvb] = Z1sum
self.Z1diag = Z1diag # this is the diagonal element of positive impedance matrix, the output of this function.
ind = 0
for pvb in self.pvbuses:
pvbusindex[pvb] = ind
ind = ind + 1
self.pvbusindex = pvbusindex
for pvb in self.pvbuses:
for pv in self.pvbuses:
d[pvb][pv] = set(self.lines[pvb]).intersection(
set(self.lines[pv]))
Z1sum = 0
for l in d[pvb][pv]:
Z1sum = Z1sum + l.Z012[1][1]
Z1[pvbusindex[pvb], pvbusindex[pv]] = Z1sum
self.d = d # this is a two dimensional dictionary. Its diagonal element contains list of lines on the path between a pv bus and root bus.
# Its of diagonal contains the common lines between path to root of two pv buses. If we sum the positive sequence impedance
# of each element, we have already the element of postive sequence impedance matrix we wish to calculate in this function.
self.Z1 = Z1
self.X1 = imag(Z1)
class rxnGraph(object):
def __init__(self, datatables, proteins=None, ligands=None, interactions=None): #, dockdf=None, seadf=None, rsimpn=None):
# Construct graph
self.graph = DiGraph()
self.data = datatables
self.proteins = proteins # list of nodes
self.ligands = ligands # list of nodes
self.protein_members = [] # subset of possible proteins, members in graph
self.protein_choices = [] # subset of possible proteins, not members in graph
self.ligand_members = [] # subset of possible ligands, members in graph
self.ligand_choices = [] # subset of possible ligands, not members in graph
self.transporter = None
self.protein_nodes = [] # moveable nodes
self.ligand_nodes = []
self.protein_allnodes = []
self.ligand_allnodes = []
self.fixed = []
self.keep = []
if proteins is not None and ligands is not None and interactions is not None:
self.construct_graph(proteins, ligands, interactions)
self.protein_members = proteins[:]
self.ligand_members = ligands[:]
def add_choices(self, protein_choices=[], ligand_choices=[]):
self.ligand_choices = ligand_choices
self.protein_choices = protein_choices
def add_transporter(self, transporter):
self.transporter = transporter
#n = self.get_start()
#self.transporter_substrate = self.graph.node[n]['label']
def copy(self):
newGraph = rxnGraph(self.data, proteins=self.proteins[:], ligands = self.ligands[:])
#dockdf=self.dockdf, seadf=self.seadf, rsimpn=self.rsimpn)
newGraph.graph = self.graph.copy()
#newGraph.constraints = self.constraints
newGraph.ligand_members = self.ligand_members[:]
newGraph.protein_members = self.protein_members[:]
newGraph.ligand_choices = self.ligand_choices[:]
newGraph.protein_choices = self.protein_choices[:]
newGraph.protein_nodes = self.protein_nodes
newGraph.ligand_nodes = self.ligand_nodes
newGraph.protein_allnodes = self.protein_allnodes
newGraph.ligand_allnodes = self.ligand_allnodes
newGraph.fixed = self.fixed
newGraph.keep = self.keep
newGraph.transporter = self.transporter
return newGraph
def construct_graph_from_path(self, proteins, ligands):
#self.graph.add_nodes_from(ligands, bipartite=0)
#self.graph.add_nodes_from(proteins, bipartite=1)
interactions = []
for i in range(len(proteins)):
interactions.append([ligands[i], proteins[i]])
if i > 0:
interactions.append([proteins[i-1], ligands[i]])
if len(ligands) == len(proteins) + 1:
interactions.append([proteins[-1], ligands[-1]])
self.construct_graph(proteins, ligands, interactions)
def construct_graph(self, proteins, ligands, interactions):
# interactions: [[protein, ligand], [protein, ligand]]
self.graph.clear()
labelposdict = {}
for i, p in enumerate(proteins):
self.graph.add_node(i, label=p, bipartite=1)
labelposdict[p] = i
self.protein_allnodes.append(i)
for j, l in enumerate(ligands):
self.graph.add_node(j + len(proteins), label=l, bipartite=0)
labelposdict[l] = j + len(proteins)
self.ligand_allnodes.append(j + len(proteins))
#for i in range(len(interactions)):
# node1 = labelposdict[interactions[i][0]]
# node2 = labelposdict[interactions[i][1]]
# self.graph.add_edge(node1, node2)
[self.graph.add_edge(labelposdict[interactions[i][0]], labelposdict[interactions[i][1]]) for i in range(len(interactions))]
self.proteins = proteins
self.ligands = ligands
self.protein_members = proteins[:]
self.ligand_members = ligands[:]
self.ligand_nodes = self.ligand_allnodes[:]
self.protein_nodes = self.protein_allnodes[:]
for i in self.fixed:
if i in self.protein_nodes:
self.protein_nodes.remove(i)
if i in self.ligand_nodes:
self.ligand_nodes.remove(i)
# updates the proteins and ligands in order
# maybe not necessary
def update_proteins_ligands(self):
n = self.get_start()
all_ligands = self.ligand_members[:]
all_ligands.extend(self.ligand_choices[:])
all_proteins = self.protein_members[:]
all_proteins.extend(self.protein_choices[:])
self.ligands = []
self.proteins = []
if n in all_proteins:
self.proteins.append(n)
elif n in all_ligands:
self.ligands.append(n)
while len(self.graph.successors(n)) > 0:
n = self.graph.successors(n)[0]
if n in all_proteins:
self.proteins.append(n)
elif n in all_ligands:
self.ligands.append(n)
# also, maybe not necessary
def get_start(self):
start = None
for n in self.graph.nodes_iter():
if len(self.graph.predecessors(n)) < 1:
start = n
if start is None:
print self.graph.edges()
return start
def get_single_predecessor_node(self, node):
pred = self.graph.predecessors(node)
if len(pred) > 0:
return pred[0]
else:
return None
def get_single_successor_node(self, node):
succ = self.graph.successors(node)
if len(succ) > 0:
return succ[0]
else:
return None
def move_remove_single(self):
all = self.protein_members[:]
all.extend(self.ligand_members[:])
node = self.choose_by_degree(1, all, 1)
if node is not None:
self.graph.remove_node(node)
if node in self.ligand_members:
self.ligand_members.remove(node)
self.ligand_choices.append(node)
elif node in self.protein_members:
self.protein_members.remove(node)
self.protein_choices.append(node)
# Removes a pair of consecutive nodes, a protein node and a ligand node
# Adds edge between predecessors to the successors of the removed nodes
def move_remove_pair(self, pnode):
lnodes = None
if self.graph.in_edges(pnode) is not None and self.graph.out_edges(pnode) is not None:
ppred = self.graph.predecessors(pnode)
psucc = self.graph.successors(pnode)
node_choices = []
if len(ppred) > 0:
node_choices.append(ppred)
if len(psucc) > 0:
node_choices.append(psucc)
if len(node_choices) > 0:
#lnodes = rand.choice(node_choices)
lnode = rand.choice(node_choices)
if len(ppred) > 0 and len(psucc) > 0:
for pp in ppred:
for ps in ppred:
self.graph.add_edge(pp, ps)
self.graph.remove_node(pnode)
self.protein_members.remove(pnode)
self.protein_choices.append(pnode)
if self.graph.in_edges(lnode) is not None and self.graph.out_edges(lnode) is not None:
lpred = self.graph.predecessors(lnode)
lsucc = self.graph.successors(lnode)
if len(lpred) > 0 and len(lsucc) > 0:
for lp in lpred:
for ls in lsucc:
self.graph.add_edge(lp, ls)
self.graph.remove_node(lnode)
self.ligand_members.remove(lnode)
self.ligand_choices.append(lnode)
# choose k nodes from a nodeset with a degree less than the limit
def choose_by_degree_in(self, upperlim_deg, nodeset, k):
choices = []
for node in nodeset:
degree = self.graph.in_degree(node)
if degree < upperlim_deg:
choices.append(node)
if len(choices) >= k:
return rand.sample(choices, k)
else:
return None
def choose_by_degree_out(self, upperlim_deg, nodeset, k):
choices = []
for node in nodeset:
degree = self.graph.out_degree(node)
if degree < upperlim_deg:
choices.append(node)
if len(choices) >= k:
return rand.sample(choices, k)
else:
return None
def choose_by_degree(self, upperlim_deg, nodeset, k):
choices = []
for node in nodeset:
degree = self.graph.degree(node)
if degree < upperlim_deg:
choices.append(node)
if len(choices) >= k:
return rand.sample(choices, k)
else:
return None
# Adds a pair of consecutive nodes to the end of the linear graph
def move_add_pair(self, pnode, lnode):
last = self.choose_by_degree_in(2, self.ligand_members, 1)[0]
if last is not None:
self.graph.add_node(pnode)
self.graph.add_node(lnode)
self.graph.add_edge(last, pnode)
self.graph.add_edge(pnode, lnode)
# Updates lists to reflect the current state of the full system
self.ligand_members.append(lnode)
self.ligand_choices.remove(lnode)
self.protein_members.append(pnode)
self.protein_choices.remove(pnode)
def move_add_edge(self):
if rand.choice([0, 1]):
pnode = self.choose_by_degree_out(2, self.protein_members, 1)
lnode = self.choose_by_degree_in(2, self.ligand_members, 1)
if pnode is not None and lnode is not None:
self.graph.add_edge(pnode[0], lnode[0])
else:
pnode = self.choose_by_degree_in(2, self.protein_members, 1)
lnode = self.choose_by_degree_out(2, self.ligand_members, 1)
if pnode is not None and lnode is not None:
self.graph.add_edge(lnode[0], pnode[0])
def move_swap_pair(self, swapLigands=True):
if swapLigands:
node1, node2 = rand.sample(self.ligand_nodes, 2)
else:
node1, node2 = rand.sample(self.protein_nodes, 2)
if (node1 not in self.fixed) and (node2 not in self.fixed):
label1 = self.graph.node[node1]['label']
label2 = self.graph.node[node2]['label']
self.graph.node[node1]['label'] = label2
self.graph.node[node2]['label'] = label1
def move_swap_neighboring_pair(self, swapLigands=True):
if swapLigands:
node = rand.choice(self.ligand_nodes)
else:
node = rand.choice(self.protein_nodes)
predprot = self.get_single_predecessor_node(node)
succprot = self.get_single_successor_node(node)
options = []
if predprot is not None:
predpred = self.get_single_predecessor_node(predprot)
if predpred is not None:
options.append(predpred)
if succprot is not None:
succsucc = self.get_single_successor_node(succprot)
if succsucc is not None:
options.append(succsucc)
if self.fixed in options:
options.remove(self.fixed)
if len(options) > 0:
swapnode = rand.choice(options)
label1 = self.graph.node[node]['label']
label2 = self.graph.node[swapnode]['label']
self.graph.node[node]['label'] = label2
self.graph.node[swapnode]['label'] = label1
# assumes linear
def move_lp_swap_pairs(self, pnode1, pnode2):
lpred = self.graph.predecessors(pnode1)
if not len(lpred) > 0:
tempnode = pnode1
pnode1 = pnode2
pnode2 = tempnode
lpred = self.graph.predecessors(pnode1)
ppred = self.graph.predecessors(lpred[0])
lsucc = self.graph.successors(pnode1)
if len(ppred) > 0:
self.graph.remove_edge(ppred[0], lpred[0])
if len(lsucc) > 0:
self.graph.remove_edge(pnode1, lsucc[0])
if len(ppred) > 0 and len(lsucc) > 0:
self.graph.add_edge(ppred[0], lsucc[0])
lsucc2 = self.graph.successors(pnode2)
if len(lsucc2) > 0:
self.graph.remove_edge(pnode2, lsucc2[0])
self.graph.add_edge(pnode1, lsucc2[0])
self.graph.add_edge(pnode2, lpred[0])
def move_swap_in(self, swapLigands=True):
if swapLigands:
select = True
while select or onodelabel in self.keep:
onode = rand.choice(self.ligand_nodes)
onodelabel = self.graph.node[onode]['label']
select = False
newnodelabel = rand.choice(self.ligand_choices)
else:
onode = rand.choice(self.protein_nodes)
onodelabel = self.graph.node[onode]['label']
newnodelabel = rand.choice(self.protein_choices)
self.graph.node[onode]['label'] = newnodelabel
if onode in self.ligand_allnodes:
self.ligand_choices.remove(newnodelabel)
self.ligand_choices.append(onodelabel)
elif onode in self.protein_allnodes:
self.protein_choices.append(onodelabel)
self.protein_choices.remove(newnodelabel)
self.protein_members.append(newnodelabel)
self.protein_members.remove(onodelabel)
# lnode is the new node to swap in
# tested, doesn't help at all
def move_swap_in_similar(self, lnode):
npdt = np.dtype([('ligand', np.str_, 35), ('tc', np.float32)])
tcarray = []
tcsum = 0.0
for lig in self.ligand_members:
tc = self.data.chsimdf.get_value(lig, lnode)
tcarray.append((lig, tc))
nptc = np.array(tcarray, dtype = npdt)
cutoff = rand.uniform(0, 0.5)
bool = nptc['tc'] > cutoff
if bool.any():
options = nptc['ligand'][bool]
else:
options = nptc['ligand']
outnode = rand.choice(options)
pred = self.get_single_predecessor_node(outnode)
succ = self.get_single_successor_node(outnode)
self.graph.add_node(lnode)
if pred is not None:
self.graph.add_edge(pred, lnode)
self.graph.remove_edge(pred, outnode)
if succ is not None:
self.graph.add_edge(lnode, succ)
self.graph.remove_edge(outnode, succ)
self.graph.remove_node(outnode)
self.ligand_members.remove(outnode)
self.ligand_choices.append(outnode)
self.ligand_members.append(lnode)
self.ligand_members.remove(lnode)
def get_subgraph_successors(self, pnode, numnodes):
currnode = pnode
subgraph_nodes = [pnode]
for n in range(numnodes - 1):
if currnode is not None:
nextnode = self.get_single_successor_node(currnode)
if nextnode is not None:
currnode = self.get_single_successor_node(nextnode)
if currnode is self.fixed:
break
if currnode is not None:
subgraph_nodes.append(currnode)
else:
return False
return subgraph_nodes
def get_subgraph_predecessors(self, pnode, numnodes):
currnode = pnode
subgraph_nodes = [pnode]
for n in range(numnodes - 1):
if currnode is not None:
nextnode = self.get_single_predecessor_node(currnode)
if nextnode is not None:
currnode = self.get_single_successor_node(nextnode)
if currnode is self.fixed:
break
if currnode is not None:
subgraph_nodes.append(currnode)
else:
return False
return subgraph_nodes.reverse()
def move_subgraph(self):
pnode = rand.choice(self.protein_nodes)
numnodes = rand.choice(range(2, len(self.protein_nodes)/2 + 1))
subgraph_nodes = self.get_subgraph_successors(pnode, numnodes)
if not subgraph_nodes:
subgraph_nodes = self.get_subgraph_predecessors(pnode, numnodes)
if subgraph_nodes:
pnode = subgraph_nodes[0]
qnode = subgraph_nodes[-1]
lnode = self.get_single_predecessor_node(pnode)
rchoices = self.protein_nodes[:]
for s in subgraph_nodes:
if s in rchoices:
rchoices.remove(s)
rnode = rand.choice(rchoices)
# Remove edges
if lnode is not None:
lpred = self.get_single_predecessor_node(lnode)
qsucc = self.get_single_successor_node(qnode)
rsucc = self.get_single_successor_node(rnode)
if lnode == rsucc:
rchoices.remove(rnode)
rnode = rand.choice(rchoices)
rsucc = self.get_single_successor_node(rnode)
if lpred is not None:
self.graph.remove_edge(lpred, lnode)
if qsucc is not None:
self.graph.add_edge(lpred, qsucc)
if qsucc is not None:
self.graph.remove_edge(qnode, qsucc)
if rsucc is not None:
self.graph.remove_edge(rnode, rsucc)
self.graph.add_edge(qnode, rsucc)
self.graph.add_edge(rnode, lnode)
def string_repr(self):
string = ''
for e in self.graph.edges():
string += "('%s', '%s'), " % (e[0], e[1])
return string.rstrip(', ')
def string_interactions(self):
string = ''
for e in self.graph.edges():
string += "'%s', '%s', " % (e[0], e[1])
return string.rstrip(', ')
def linear_to_string(self):
string = ''
if self.transporter is not None and self.transporter in self.data.dockdf.index:
string += '%s -> ' % self.transporter
n = self.get_start()
string += self.graph.node[n]['label']
while len(self.graph.successors(n)) > 0:
n = self.graph.successors(n)[0]
label = self.graph.node[n]['label']
string += ' -> '
string += label
return string
def add_constraint(self, fun):
self.constraints.append(fun)
def test_constraints(self):
for fun in self.constraints:
if not fun(self):
return False
return True
def _construct(self):
"""
We want to build the type of DFG that's used in "Automated Ident. of Crypto
Primitives in Binary Code with Data Flow Graph Isomorphisms." Unlike that
paper, however, we're building it on Vex IR instead of assembly instructions.
"""
cfg = self._cfg
p = self.project
dfgs = {}
l.debug("Building Vex DFG...")
for node in cfg.nodes():
try:
if node.simprocedure_name == None:
irsb = p.factory.block(node.addr).vex
else:
l.debug("Cannot process SimProcedures, ignoring %s" % node.simprocedure_name)
continue
except Exception as e:
l.debug(e)
continue
tmpsnodes = {}
storesnodes = {}
putsnodes = {}
statements = irsb.statements
dfg = DiGraph()
for stmt_idx, stmt in enumerate(statements):
# We want to skip over certain types, such as Imarks
if self._need_to_ignore(node.addr, stmt, stmt_idx):
continue
# break statement down into sub-expressions
exprs = stmt.expressions
stmt_node = stmt
dfg.add_node(stmt)
if stmt.tag == 'Ist_WrTmp':
tmpsnodes[stmt.tmp] = stmt_node
if exprs[0].tag == 'Iex_Binop':
if exprs[1].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[1].tmp], stmt_node)
else:
dfg.add_edge(exprs[1], stmt_node)
if exprs[2].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[2].tmp], stmt_node)
else:
dfg.add_edge(exprs[2], stmt_node)
elif exprs[0].tag == 'Iex_Unop':
dfg.remove_node(stmt_node)
if exprs[1].tag == 'Iex_RdTmp':
tmpsnodes[stmt.tmp] = copy(tmpsnodes[exprs[1].tmp])
tmpsnodes[stmt.tmp].tmp = stmt.tmp
else:
tmpsnodes[stmt.tmp] = exprs[1]
elif exprs[0].tag == 'Iex_RdTmp':
tmpsnodes[stmt.tmp] = copy(tmpsnodes[exprs[0].tmp])
tmpsnodes[stmt.tmp].tmp = stmt.tmp
elif exprs[0].tag == 'Iex_Get':
if putsnodes.has_key(exprs[0].offset):
dfg.add_edge(putsnodes[exprs[0].offset], stmt_node)
if len(exprs) > 1 and exprs[1].tag == "Iex_RdTmp":
dfg.add_edge(tmpsnodes[exprs[1].tmp], stmt_node)
elif len(exprs) > 1:
dfg.add_edge(exprs[1], stmt_node)
elif exprs[0].tag == 'Iex_Load':
if exprs[1].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[1].tmp], stmt_node)
else:
dfg.add_edge(exprs[1], stmt_node)
else:
# Take a guess by assuming exprs[0] is the op and any other expressions are args
for e in exprs[1:]:
if e.tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[e.tmp], stmt_node)
else:
dfg.add_edge(e, stmt_node)
elif stmt.tag == 'Ist_Store':
if exprs[0].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[0].tmp], stmt_node)
elif exprs[0].tag == 'Iex_Const':
dfg.add_edge(exprs[0], stmt_node)
if exprs[1].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[1].tmp], stmt_node)
else:
dfg.add_edge(exprs[1], stmt_node)
elif stmt.tag == 'Ist_Put':
if exprs[0].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[0].tmp], stmt_node)
elif exprs[0].tag == 'Iex_Const':
dfg.add_edge(exprs[0], stmt_node)
putsnodes[stmt.offset] = stmt_node
elif stmt.tag == 'Ist_Exit':
if exprs[0].tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[exprs[0].tmp], stmt_node)
elif stmt.tag == 'Ist_Dirty':
tmpsnodes[stmt.tmp] = stmt_node
elif stmt.tag == 'Ist_CAS':
tmpsnodes[stmt.oldLo] = stmt_node
else:
for e in stmt.expressions:
if e.tag == 'Iex_RdTmp':
dfg.add_edge(tmpsnodes[e.tmp], stmt_node)
else:
dfg.add_edge(e, stmt_node)
for vtx in dfg.nodes():
if dfg.degree(vtx) == 0:
dfg.remove_node(vtx)
if dfg.size() > 0:
dfgs[node.addr] = dfg
return dfgs
