def add_ip_tables(self, host=None): """Add ip table rules. Accepts incoming traffic on ports. Redirects traffic on simulation ports to default ports. Configuration specified in server config. """ flog.info("Accepting and redirecting ports (host=%s)" % host) if host: hostname = os.environ.get("FILTER_HOST", "filter") filter_host = socket.gethostbyname(hostname) assert host, "Unable to get host from hostname %s" % hostname flog.debug("Redirect connections to host %s to %s" % (filter_host, host)) self.run( "-t nat -A POSTROUTING -o eth0 -j SNAT --to-source %s" % filter_host, tool=Command.Iptables, ) for protocol, protocol_set in self.server_config["ports"].items(): for default_port, port_set in protocol_set.items(): if default_port != "null": port_set.append(default_port) if not self.accept_ports(protocol=protocol, ports=port_set): return False if host: default_port = host if default_port != "null" and not self.redirect_ports( protocol=protocol, src_list=port_set, dst=default_port): return False self.show_ip_tables() return True
def send(self, data): """Send encoded data to server.""" try: flog.info(f"Sending data of length: {len(data)}") self.client.send(data) except (BlockingIOError, ConnectionResetError, socket.timeout) as e: flog.debug(e)
def tcp_handle(client_sock, killable=False): """Echo data over TCP socket.""" flog.info("Launch TCP Handler.") while True: try: data = client_sock.recv(max_buffer) try: data_str = data.decode("utf-8") except Exception: pass if killable and data_str and "kill" in data_str: flog.debug("Killing Socket...") pid = os.getpid() os.kill(pid, signal.SIGKILL) if data: flog.debug("Received TCP data: %s" % data) client_sock.send(data) else: client_sock.close() break except Exception as ex: flog.warning("TCP Handler Exception: {}".format(ex)) client_sock.close() break flog.info("End TCP Handler")
def check_file(self, size): """Check if file exists.""" file_path = self._generate_file_path(size) flog.info("Checking if {} exists".format(file_path)) if not os.path.exists(file_path): flog.warning("{} does not exist".format(file_path)) return False flog.info("{} exists".format(file_path)) return True
def do_POST(self): """Start TCP client from request.""" global DEFAULT_TIMEOUT, SERVICES, LOCAL_PORT_UDP bad_request_str = b"<html><body><h1>Bad Request!</h1></body></html>" try: form = cgi.FieldStorage(fp=self.rfile, headers=self.headers, environ={"REQUEST_METHOD": "POST"}) service = form.getvalue("service", None) if not service or service not in SERVICES: self._set_headers(400) self.wfile.write(bad_request_str) return request_vals = { "address": form.getvalue("address", None), "port": form.getvalue("port", None), "message": form.getvalue("message", "Hello from Flake Client!"), "echo": form.getvalue("echo", False), } flog.info(request_vals) if service in [ "tcp_client", "udp_client" ] and not (request_vals["address"] and request_vals["port"] and request_vals["port"].isdigit() and (request_vals["message"] or request_vals["echo"])): self._set_headers(400) self.wfile.write(bad_request_str) return local_port = None if service == "udp_client": LOCAL_PORT_UDP += 1 local_port = LOCAL_PORT_UDP if service in ["tcp_client", "udp_client"]: echo_client = EchoClient( address=request_vals["address"], port=int(request_vals["port"]), message=request_vals["message"], echo=request_vals["echo"], mode="TCP" if service == "tcp_client" else "UDP", local_port=local_port, timeout=DEFAULT_TIMEOUT, ) flog.info(f"Starting {service} thread.") _thread.start_new_thread(echo_client.run, ()) self._set_headers(200) self.wfile.write( b"<html><body><h1>Request Received, starting %s!</h1></body></html>" % service.encode()) except Exception as e: self._set_headers(400) self.wfile.write( b"<html><body><h1>Bad Request!</h1><br><p>%s</p></body></html>" % repr(e).encode("utf-8"))
def create_file(self, size): """Create data file (image).""" file_path = self._generate_file_path(size) flog.info("Creating a image file %s of size %s" % (file_path, size)) byte_size = self._parse_size(size) # Need a slow growing series sqrt(x)/2 works for this. size = int(np.sqrt(byte_size) / 2) d = np.random.random_sample((size, size)).astype(np.float32) im = Image.fromarray(d, mode="F") im.save(file_path)
def check_message(self, message): "Check if message has command encoded." try: data_str = message.decode("utf-8") flog.debug(data_str) if "close" in data_str: flog.info("Received close command") return "close" except Exception: pass return ""
def create_file(self, size): """Create data file (audio).""" file_path = self._generate_file_path(size) flog.info("Creating a audio file %s of size %s" % (file_path, size)) byte_size = self._parse_size(size) # Need to take fraction of the size to get correct output size. length = byte_size / 352830 sample_rate = 44100 frequency = 440.0 t = np.linspace(0, length, int(sample_rate * length)) y = np.sin(frequency * 2 * np.pi * t) # Has frequency of 440Hz wavfile.write(file_path, sample_rate, y)
def get_file(self, path): """Parse path and return file.""" for file_type in self.config.keys(): if file_type in path: path = os.path.join(self.root, "files", path) break else: path = os.path.join(self.root, path) if not os.path.exists(path): return str.encode("Incorrect path {}".format(path)) with open(path, "rb") as f: flog.info("Sending file {}".format(path)) return f.read()
def create_file(self, size): """Create data file (video).""" file_path = self._generate_file_path(size) flog.info("Creating a video file %s of size %s" % (file_path, size)) byte_size = self._parse_size(size) file_str = self.template template_size = self._parse_size("0.5MB") while template_size < byte_size: file_str += "|{}".format(self.template) byte_size -= template_size cmd = 'ffmpeg -i "concat:{}" -c copy {}'.format(file_str, file_path) assert os.system(cmd) == 0, "Failed to create video"
def create_file(self, size): """Create data file (binary).""" file_path = self._generate_file_path(size) flog.info("Creating a binary file %s of size %s" % (file_path, size)) byte_size = self._parse_size(size) chunk_size = 1024 * 1024 with open(file_path, "wb") as f: while byte_size > 0: write_size = byte_size if write_size > chunk_size: write_size = chunk_size f.write(os.urandom(write_size)) byte_size -= write_size
def run(self): """Run pcap.""" flog.info("Creating pcap listener process") self.create_listener_process() while self.listener_process.poll() is None: out = "" # Scan listener packet input for 1 second with funtion_timeout(1): out = self.listener_process.stdout.readline() if out: if not self.pcap_started: flog.info("Capture started...") log_file = self.get_log_file() self.pcap_started = True # Reset time self.last_received_packet = time.time() # Check if timeout reached between packet inputs if self.pcap_started and is_timeout(self.last_received_packet, self.timeout): flog.info("Capture done.") self.kill_listener() self.clear_vars() flog.info("Saving pcap to log file: {}".format(log_file)) self.save_pcap(log_file) self.create_listener_process() flog.warning(self.listener_process.communicate())
def udp_handle(self): """Echo data over UDP socket.""" flog.info("Launch UDP Handler.") while True: data, addr = self.udp_server.recvfrom(max_buffer) if data: if data == b"0 byte test": flog.debug("Received 0 byte test") self.udp_server.sendto(b"", addr) else: flog.debug("Received UDP data: %s" % data) self.udp_server.sendto(data, addr) if not data: self.udp_server.close()
def run_tcp_tls_server(): """Run TCP TLS server.""" flog.info("Starting TCP TLS echo server") config = ServerConfig() tcp_TLS_server = TcpTLSServer(config) tcp_TLS_server_mutual = TcpTLSServer(config, mutual=True) flog.debug("Launch mutual authentication server thread.") mutual_thread = threading.Thread(target=tcp_TLS_server_mutual.run, args=()) mutual_thread.start() flog.debug("Launch server authentication server thread.") server_thread = threading.Thread(target=tcp_TLS_server.run, args=()) server_thread.start() mutual_thread.join() server_thread.join()
def create_file(self, size): """Create data file (text).""" file_path = self._generate_file_path(size) flog.info("Creating a text file %s of size %s" % (file_path, size)) byte_size = self._parse_size(size) chunk_size = 1024 with open(file_path, "w") as f: while byte_size > 0: write_size = byte_size if write_size > chunk_size: write_size = chunk_size f.write("".join( np.random.choice( list(string.ascii_letters + string.digits), write_size))) byte_size -= write_size
def run(self): """Run TCP/UDP client.""" flog.info(f"Connecting to {self.mode} server at {self.address}:{self.port}") flog.debug(f"Echo: {self.echo}") if not self.echo: flog.debug(f"Message Size: {len(self.message)}") if self.mode == "TCP": self.run_tcp() elif self.mode == "UDP": self.run_udp() else: flog.error( f"Cannot run client in mode: {self.mode}\nplease use either TCP or UDP" ) return flog.info(f"End {self.mode} Client")
def configure_files(force=False, config_file=CONFIG): """Configure server files.""" config = ConfigHandler(config_file) configuration_times = {} total_time = time.time() for data_type, data_values in config.items(): start_time = time.time() data_handler = DataFile.get_data_handler(data_type=data_type, config=data_values, server_config=config.server) data_handler.configure(force) configuration_times[data_type] = "%ss" % format( time.time() - start_time, ".2f") configuration_times["Total"] = "%ss" % format(time.time() - total_time, ".2f") flog.info("Time to configure files:\n%s" % configuration_times) return True
def configure(self, with_filtering=True, with_services=True): """Configure server based on config rules.""" host = None if not with_services: hostname = os.environ.get("SERVICES_HOST", "server") host = socket.gethostbyname(hostname) assert host, "Unable to get host from hostname %s" % hostname assert self.add_ip_tables(host), "Failed to add to ip tables" if not with_filtering: return # Reset current rules self.clear_htb() assert self._create_htb(), "Failed to create htb" for name, setup in self.config.items(): flog.info("Configuring for set: {name}".format(name=name)) flog.info( "Packet delays: {delay}, Packet loss: {loss}, Ports: {ports}". format( delay=setup["packet_delay"], loss=setup["packet_loss"], ports=setup["ports"], )) for delay, ports in zip(setup["packet_delay"], setup["ports"]): self.index += 1 if not with_services: # If the services are remote, the delay is also applied when # we forward the packet, so we need to divide it in half. delay = "%fms" % (float(delay.strip("ms")) / 2) assert self.add_egress_rule( port=ports, class_id=self.index, packet_loss=setup["packet_loss"], packet_delay=delay, ), "Failed to set egress rule" # If the services are 'local', we also need to add an ingress rule. # With the services are remote, egress rule also applies to ingress traffic # when it gets forwarded to remote port. if with_services: assert self.add_ingress_rule( port=ports, packet_loss=setup["packet_loss"] ), "Failed to set ingress rule"
def tcp_tls_handle(conn_stream, client_sock): """Echo data over TLS connection.""" flog.info("TCP TLS Handle Launched.") while True: try: data = conn_stream.read() if data: flog.debug("Received TCP TLS: {}".format(data)) conn_stream.send(data) if not data: conn_stream.close() client_sock.close() break except Exception as ex: flog.warning("TCP TLS Handler Exception: {}".format(ex)) conn_stream.close() client_sock.close() break flog.info("End TCP TLS Handler.")
def main(): """Parse arguments.""" parser = argparse.ArgumentParser(description="Flake Legato CLI Tool.") parser.add_argument( "action", type=str, help="action for flake server", choices=["configure", "test"], ) parser.add_argument("-f", "--force", action="store_true", help="force action") args = parser.parse_args() server_manager = Server() flog.info("Server: filtering[%s] services[%s]" % (server_manager.with_filtering, server_manager.with_services)) return server_manager.run_action(action=args.action, force=args.force)
def receive(self): """Receive encoded data from server.""" data = b"" wait_time = 1 got_data = False for _ in range(0, self.timeout, wait_time): sockets, _, _ = select.select((self.client,), (), (), wait_time) for sock in sockets: data += sock.recv(MAX_BUFFER) got_data = data != b"" flog.debug(f"Data part received length: {len(data)}") break else: if got_data: flog.debug("Data parts complete") break if not got_data: return None flog.info(f"Received message of length: {len(data)} bytes") return data
def run(self): """Run TCP TLS server.""" TLS_context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2) flog.debug("Loading cert and keyfile") TLS_context.load_cert_chain(certfile=self.fullchain, keyfile=self.privkey) if self.is_mutual is True: flog.debug("Setting server to mutual authentication.") # CA root used for verifying Client certificates TLS_context.load_verify_locations(cafile=self.CAroot) TLS_context.verify_mode = ssl.CERT_REQUIRED self.tcp_server.bind((self.local_ip, self.port)) self.tcp_server.listen(5) flog.debug("Starting TCP TLS server on {}".format(self.port)) while True: # listen for tcp then wrap with TLS client_sock, addr = self.tcp_server.accept() client_sock.settimeout(self.timeout) flog.info("TCP Socket Connected.") try: conn_stream = TLS_context.wrap_socket(client_sock, server_side=True) TLS_thread = threading.Thread(target=tcp_tls_handle, args=(conn_stream, client_sock)) TLS_thread.start() except Exception as ex: flog.info("Could not wrap socket. Closing TCP Socket.") flog.info(ex) client_sock.close()
def run_udp(self): """Run UDP client.""" self.client = socket(AF_INET, SOCK_DGRAM) self.client.settimeout(self.timeout) data = self.message.encode() addr = (self.address, self.port) self.client.bind(("0.0.0.0", self.local_port)) while True: if not self.echo: data = self.message.encode() try: flog.info(f"Sending data of length: {len(data)}") self.client.sendto(data, addr) data, addr = self.client.recvfrom(MAX_BUFFER) flog.info(f"Received message of length: {len(data)} bytes") if not data or self.check_message(data) == "close": self.client.close() break except Exception as ex: flog.warning(f"UDP Client Exception: {ex}") self.client.close() break
def show_ip_tables(self): """Show ip tables rules.""" flog.info("IP Table Rules:") flog.info( self.run("-n -L --line-numbers", tool=Command.Iptables, output=True)) flog.info( self.run("-t nat -n -L --line-numbers", tool=Command.Iptables, output=True))
def run_echo_server(): """Run echo servers.""" flog.info("Starting TCP/UDP Echo server") config = ServerConfig() echo_servers = EchoServer(config) echo_servers.run()
def configure(self, force=False): """Configure server at startup.""" if self.with_services: flog.info("---- Configuring Flake ----") assert self.add_files(force), "Failed to add files." flog.info("---- Successfully added files ----") assert port_publisher.publish_ports(), "Failed to publish ports." flog.info("---- Successfully published ports ----") assert traffic_manager.configure_server_rules( with_filtering=self.with_filtering, with_services=self.with_services ), "failed to configure traffic rules." flog.info("---- Successfully configured traffic rules ----") if self.with_services: assert self.start_udp_server(), "failed to start UDP server." flog.info("---- Successfully started UDP server ----") assert self.configure_iperf(), "failed to configure iperf server." flog.info("---- Successfully configured iperf on server ----") assert self.start_echo_servers( ), "failed to start TCP/UDP Echo servers." flog.info("---- Successfully started TCP/UDP Echo servers ----") assert self.start_tcp_tls_server( ), "failed to start TCP TLS server." flog.info("---- Successfully started TCP TLS Echo server ----") assert (self.start_dynamic_http_server() ), "failed to start dynamic http server." flog.info("---- Successfully started http TCP client server ----") assert self.configure_pcap(), "failed to start pcap." flog.info("---- Successfully started pcap on server ----") assert self.configure_cron(), "failed to start cron." flog.info("---- Successfully started cron on server ----") return True
def main(): """Run Http TCP server.""" flog.info("Starting Server") config = ServerConfig() server = HttpRequestServer(config=config) server.run()
def run(self): """Run http server.""" httpd = HTTPServer(self.address, RequestHandler) flog.info("Server running at localhost:%d..." % self.port) httpd.serve_forever()
def run_server(config_file=CONFIG): """Run udp server.""" flog.info("Starting UDP server") config = ConfigHandler(config_file) udp_server = UdpServer(config) udp_server.run()
def start_iperf(config_file=CONFIG): """Start iperf on server.""" flog.info("Starting iperf") config = ConfigHandler(config_file) iperf_handler = Iperf(config) return iperf_handler.start()