diff --git a/start.py b/start.py index fc22fb7..0ae808d 100755 --- a/start.py +++ b/start.py @@ -20,6 +20,12 @@ def puts(text, *args, color=colors.white, is_bold=False, **kwargs): def sep(): puts("-----------------------------------", is_bold=True) +def composecmd(cmd): + return os.system(f"(hash docker-compose &> /dev/null && docker-compose -p firegex {cmd} || exit 0) || (hash docker &> /dev/null && docker compose -p firegex {cmd} || exit 0) || echo 'Docker not found!, please install docker and docker-compose'") + +def dockercmd(cmd): + return os.system(f"(hash docker &> /dev/null && docker {cmd} || exit 0) || echo 'Docker not found!, please install docker and docker-compose'") + parser = argparse.ArgumentParser() parser.add_argument('--port', "-p", type=int, required=False, help='Port where open the web service of the firewall', default=4444) parser.add_argument('--threads', "-t", type=int, required=False, help='Number of threads started for each service/utility', default=-1) @@ -106,16 +112,16 @@ if not args.no_autostart: try: if args.restart: puts("Running 'docker-compose restart'\n", color=colors.green) - os.system("docker-compose -p firegex restart 2> /dev/null || docker compose -p firegex restart") + composecmd("restart") elif args.stop: puts("Running 'docker-compose down'\n", color=colors.green) - os.system("docker-compose -p firegex down 2> /dev/null || docker compose -p firegex restart") + composecmd("down") else: if not args.build: puts("Downloading docker image from github packages 'docker pull ghcr.io/pwnzer0tt1/firegex'", color=colors.green) - os.system("docker pull ghcr.io/pwnzer0tt1/firegex") + dockercmd("pull ghcr.io/pwnzer0tt1/firegex") puts("Running 'docker-compose up -d --build'\n", color=colors.green) - os.system("docker-compose -p firegex up -d --build 2> /dev/null || docker compose -p firegex up -d --build") + composecmd("up -d --build") finally: if not args.keep: os.remove("docker-compose.yml") diff --git a/tests/ph_test.py b/tests/ph_test.py index 66277a3..02cf222 100755 --- a/tests/ph_test.py +++ b/tests/ph_test.py @@ -49,7 +49,7 @@ time.sleep(0.5) def checkData(should_work): res = None try: res = server.sendCheckData(secrets.token_bytes(432)) - except ConnectionRefusedError: res = False + except (ConnectionRefusedError, TimeoutError): res = None if res: if should_work: puts(f"Successfully received data ✔", color=colors.green) else: puts("Test Failed: Connection wasn't blocked ✗", color=colors.red); exit_test(1) diff --git a/tests/utils/tcpserver.py b/tests/utils/tcpserver.py index 0cfd37d..4e31833 100644 --- a/tests/utils/tcpserver.py +++ b/tests/utils/tcpserver.py @@ -26,7 +26,8 @@ class TcpServer: def sendCheckData(self,data): s = socket.socket(socket.AF_INET6 if self.ipv6 else socket.AF_INET, socket.SOCK_STREAM) - s.connect(('::1' if self.ipv6 else '127.0.0.1', self.proxy_port if self.proxy_port else self.port)) + s.settimeout(1) + s.connect(('::1' if self.ipv6 else '127.0.0.1', self.proxy_port if self.proxy_port else self.port), ) s.sendall(data) received_data = s.recv(4096) s.close()