rewrite tcp port check

This commit is contained in:
kiriharu 2021-01-04 23:13:51 +03:00
parent 16343d6ea0
commit cd2e0002f8
3 changed files with 50 additions and 37 deletions

View File

@ -26,9 +26,9 @@ def http_check():
@access_token_required
def tcp_port_check():
target = request.args.get("target", None)
port = int(request.args.get("port", 80))
port = int(request.args.get("port", None))
if not target:
if not target or not port:
abort(400)
checker = TCPPortChecker(target, port)

View File

@ -0,0 +1,42 @@
from core.coretypes import Response, PortResponse, ResponseStatus, ErrorPayload, ErrorCodes
from .base import BaseChecker
import socket
class PortChecker(BaseChecker):
def __init__(self, target: str, port: int, sock: socket.socket):
self.port = port
self.sock = sock
super().__init__(target)
def check(self) -> Response:
# 2 seconds timeout...
self.sock.settimeout(2)
try:
res = self.sock.connect_ex((self.target, self.port))
except socket.gaierror:
return Response(
status=ResponseStatus.ERROR,
payload=ErrorPayload(
message="Invalid hostname",
code=ErrorCodes.InvalidHostname
),
node=self.node_info
)
if res == 0:
self.sock.close()
return Response(
status=ResponseStatus.OK,
payload=PortResponse(open=True),
node=self.node_info
)
self.sock.close()
return Response(
status=ResponseStatus.OK,
payload=PortResponse(open=False),
node=self.node_info
)

View File

@ -1,42 +1,13 @@
from core.coretypes import Response, PortResponse, ResponseStatus, ErrorPayload, ErrorCodes
from .base import BaseChecker
from .port import PortChecker
import socket
class TCPPortChecker(BaseChecker):
class TCPPortChecker(PortChecker):
def __init__(self, target: str, port: int):
self.port = port
super().__init__(target)
def check(self) -> Response:
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
# 2 seconds timeout...
sock.settimeout(2)
try:
res = sock.connect_ex((self.target, self.port))
except socket.gaierror:
return Response(
status=ResponseStatus.ERROR,
payload=ErrorPayload(
message="Invalid hostname",
code=ErrorCodes.InvalidHostname
),
node=self.node_info
)
if res == 0:
sock.close()
return Response(
status=ResponseStatus.OK,
payload=PortResponse(open=True),
node=self.node_info
)
sock.close()
return Response(
status=ResponseStatus.OK,
payload=PortResponse(open=False),
node=self.node_info
super().__init__(
target=target,
port=port,
sock=socket.socket(socket.AF_INET, socket.SOCK_STREAM)
)