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 @access_token_required
def tcp_port_check(): def tcp_port_check():
target = request.args.get("target", None) 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) abort(400)
checker = TCPPortChecker(target, port) 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 .port import PortChecker
from .base import BaseChecker
import socket import socket
class TCPPortChecker(BaseChecker): class TCPPortChecker(PortChecker):
def __init__(self, target: str, port: int): def __init__(self, target: str, port: int):
self.port = port super().__init__(
super().__init__(target) target=target,
port=port,
def check(self) -> Response:
sock=socket.socket(socket.AF_INET, socket.SOCK_STREAM) 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
) )