class TestServerAndComponent(TestCase):
@classmethod
def setUpClass(self):
class MockComponent(Component):
def process(self, json):
return {"test": "test"}
self.component = MockComponent()
self.server = Server("test", "4343")
self.server.add_component_post("/test1", self.component)
self.server.add_component_get("/test2", self.component)
self.process = Process(target=self.server.run)
self.process.start()
time.sleep(5)
@classmethod
def tearDownClass(self):
os.system("kill " + str(self.process.pid))
time.sleep(5)
def test_component_post_method_called(self):
full_url = "http://127.0.0.1:4343/test1"
request = json.loads('{"test":"test"}')
r = requests.post(url=full_url, json=request)
self.assertTrue("test" in r.text)
def test_component_get_method_called(self):
full_url = "http://127.0.0.1:4343/test2"
request = json.loads("{}")
r = requests.get(url=full_url, json=request)
self.assertTrue("test" in r.text)
def __init__(self):
self.server = Server("database", config["database"])
self.database = MemoryDB()
self.server.add_component_post("/vulnerability",
DBQuery(self.database))
self.server.add_component_post("/privileges",
DBPrivileges(self.database))
def setUpClass(self):
class MockComponent(Component):
def process(self, json):
return {"test": "test"}
self.component = MockComponent()
self.server = Server("test", "4343")
self.server.add_component_post("/test1", self.component)
self.server.add_component_get("/test2", self.component)
self.process = Process(target=self.server.run)
self.process.start()
time.sleep(5)
def __init__(self,
graph,
sniffer_client=LocalClient(config["sniffer"]),
batch_threads=1):
self.graph = graph
self.populator = Populator(self.graph, batch_threads=batch_threads)
self.server = Server("graph", config["graph"])
self.server.add_component_get("/graph", GraphExporter(graph))
self.server.add_component_post("/merge", GraphMerge(graph))
self.sniffer_client = sniffer_client
def __init__(self, device):
shared_lock = Lock()
shared_list = []
self.server = Server("sniffer", config["sniffer"])
self.server.add_component_get("/newpackets",
PacketExporter(shared_list, shared_lock))
if device is None:
self.daemon = SniffingDaemon(shared_list, shared_lock)
else:
self.daemon = SniffingDaemon(shared_list, shared_lock, connections=open_connection(device))
class DatabaseService():
"""
The DatabaseService encapsulates all the dependecies of the
database service:
* the memorydb
* the server
"""
def __init__(self):
self.server = Server("database", config["database"])
self.database = MemoryDB()
self.server.add_component_post("/vulnerability",
DBQuery(self.database))
self.server.add_component_post("/privileges",
DBPrivileges(self.database))
def __init__(self, slave_port, master_ip, master_port, client_cls=Client):
self.slave_port = slave_port
self.slave_ip = get_host_ip()
self.client_cls = client_cls
self.master_client = self.client_cls("http://" + master_ip,
master_port)
self.membership_list = []
self.server = Server("slave", slave_port)
self.server.add_component_get("/healty", HealthCheck())
self.server.add_component_post("/membership", SlaveMembership(self))
self.server.add_component_post("/multicast", MessageReceiver(self))
self.dissemination_constant = 5
self.graph_sharing = GraphSharing()
def main():
args = sys.argv
schedule_recurring = False
if '-s' in args:
args.remove('-s')
schedule_recurring = True
keywords = args[1:]
parser = ParseDG(keywords)
mail = Mail(config)
try:
with Server(config) as server:
def update_and_send():
parser.run()
results = parser.results
mail.make_sender(server, results)
mail.send()
print "Update sent!"
if schedule_recurring:
minutes = config.TIME
print "Email update every {} minute(s)".format(minutes)
update_and_send()
schedule.every(minutes).minutes.do(update_and_send)
while True:
schedule.run_pending()
time.sleep(1)
else:
update_and_send()
except Exception as e:
print e
print "Closing Service"
class GraphService():
"""
The GraphService encapsulates all the dependecies of the
graph service:
* a local client to the sniffer
* a server
* a populator
"""
def __init__(self,
graph,
sniffer_client=LocalClient(config["sniffer"]),
batch_threads=1):
self.graph = graph
self.populator = Populator(self.graph, batch_threads=batch_threads)
self.server = Server("graph", config["graph"])
self.server.add_component_get("/graph", GraphExporter(graph))
self.server.add_component_post("/merge", GraphMerge(graph))
self.sniffer_client = sniffer_client
def get_edge(self, packet):
"""
Returns an edges from a packet dictionary. The packet must have
fields "src" and "dest".
"""
src, dest = packet["src"], packet["dest"]
return Node(src), Node(dest)
def update(self):
"""
Updates a graph with new packets requested from the sniffer
microservice.
"""
new_packets = self.sniffer_client.get("/newpackets", default=[])
self.graph.lock.acquire()
for packet in new_packets:
edge = self.get_edge(packet)
if edge is not None:
self.graph.add_edge(*edge)
self.graph.lock.release()
class SniffingService():
"""
Wrapper over:
- the server with a PacketExporter component
- the shared_list of packets and the shared_lock
Shares the `shared_list` of packets with the sniffing daemon.
"""
def __init__(self, device):
shared_lock = Lock()
shared_list = []
self.server = Server("sniffer", config["sniffer"])
self.server.add_component_get("/newpackets",
PacketExporter(shared_list, shared_lock))
if device is None:
self.daemon = SniffingDaemon(shared_list, shared_lock)
else:
self.daemon = SniffingDaemon(shared_list, shared_lock, connections=open_connection(device))
def init_server(ipmi_user, ipmi_addr, ipmi_pass, mac_addr):
# sql = "SELECT ID FROM " + 'SERVERS'
# cur.execute(sql)
# if cur = "":
# id=
id = "00000000-0000-0000-0000-000000000001"
name = "server1"
ipv4_addr = "172.16.166.34"
server = Server(id, name, ipmi_user, ipmi_addr, ipmi_pass, mac_addr,
ipv4_addr)
return server
class Master():
def __init__(self):
self.membership_list = []
self.server = Server("master", MASTER_DEFAULT_PORT)
self.server.add_component_post("/register", MasterReceive(self))
def register(self, registration, client_cls=Client):
"""
Register slave in membership list.
"""
logger.info("Received register {}:{}".format(registration["ip"],
registration["port"]))
# The servers are single threaded, one-connection at a time
client = client_cls("http://" + registration["ip"],
registration["port"])
self.membership_list.append(client)
def broadcast(self):
"""
Create a JSON with information about each client, i.e. slave
in the membership list and send the broadcast data to all
other slaves.
"""
logger.info(colored.yellow("Broadcasting membership...."))
broadcast = {
"members": [{
"ip": client.url.split("/")[2],
"port": client.port,
} for client in self.membership_list]
}
# Send the broadcast to each client
for client in self.membership_list:
client.get("/healty")
client.post("/membership", broadcast)
def inference_service(env=os.environ.copy()):
"""
The inference service provides an overlay over the Mulval service.
"""
client = LocalClient(config["graph"])
server = Server("inference", config["inference"])
server.add_component_get("/attack_graph", AttackGraphExporter(client, env))
server.run()
# -*- coding: utf-8 -*-
from multiprocessing import Process
import pytest
from service.server import Server
from setting import Setting
from utils.ope_run_data import OpeRunData
def run_test_action(index):
Setting.UDID = OpeRunData().get_value(f'device_{index}', 'udid')
Setting.PORT = OpeRunData().get_value(f'device_{index}', 'port')
Setting.Account = Setting.zpjAccount #选择账号(多设备同账号 多设备多账号 IF判断)
pytest.main(['-s', '-q', '--alluredir', './report/xml', 'test_app/'])
if __name__ == '__main__':
s = Server()
s.start_appium_server()
for i in range(len(s.device_list)):
p = Process(target=run_test_action, args=(i, ))
p.start()
p.join()
def __init__(self):
self.membership_list = []
self.server = Server("master", MASTER_DEFAULT_PORT)
self.server.add_component_post("/register", MasterReceive(self))
class Slave():
def __init__(self, slave_port, master_ip, master_port, client_cls=Client):
self.slave_port = slave_port
self.slave_ip = get_host_ip()
self.client_cls = client_cls
self.master_client = self.client_cls("http://" + master_ip,
master_port)
self.membership_list = []
self.server = Server("slave", slave_port)
self.server.add_component_get("/healty", HealthCheck())
self.server.add_component_post("/membership", SlaveMembership(self))
self.server.add_component_post("/multicast", MessageReceiver(self))
self.dissemination_constant = 5
self.graph_sharing = GraphSharing()
def join(self):
"""Send request to the master for addition to membership list.
Slave sends its IP and port number."""
logger.info(colored.yellow("Slave requesting join."))
self.master_client.post("/register", {
"ip": self.slave_ip,
"port": self.server.port
})
def update_membership(self, membership_list):
"""Epidemic style dissemination to slave neighbours."""
logger.info(colored.green("Updating membership...."))
if "members" not in membership_list:
return
membership_list = membership_list["members"]
self.membership_list = []
for member in membership_list:
if self.slave_ip == member["ip"] and self.slave_port == member[
"port"]:
continue
client = self.client_cls("http://" + member["ip"], member["port"])
self.membership_list.append(client)
logger.info("Membership list updated: {} members.".format(
len(self.membership_list)))
def get_current_broadcast(self):
"""Returns the membership list for this slave."""
return self.membership_list
def get_current_multicast(self):
"""Returns the list of slaves this slave can send messages to.
This is the randomized part that can be proven to be O(log n) for
the whole dissemination."""
multicast_list = list(self.membership_list[:])
random.shuffle(multicast_list)
return multicast_list[:self.dissemination_constant]
def disseminate(self, multicast_list, message):
"""Send message to all multicast nodes this slave can communicate to."""
logger.info("Running dissemination.")
for client in multicast_list:
logger.info("Sending message to {}:{}.".format(
client.url, client.port))
client.post("/multicast", message)
def run(self):
"""Slave main function loop."""
self.join()
time.sleep(5)
while True:
time.sleep(10)
multicast_list = self.get_current_multicast()
multicast_message = {
"ip": self.slave_ip,
"port": self.slave_port,
"graph": self.graph_sharing.snapshoot()
}
self.disseminate(multicast_list, multicast_message)