51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
import socket
|
|
import struct
|
|
import time
|
|
|
|
from colony.core.logger.logging import get_logger
|
|
|
|
|
|
class Discover:
|
|
DISCOVER_PORT = 28149
|
|
MAGIC = b"OK"
|
|
CLIENT_TIMEOUT = 2
|
|
|
|
@staticmethod
|
|
def listen(real_port, host="0.0.0.0", discover_port=None):
|
|
logger = get_logger(Discover.__name__)
|
|
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
s.bind((host, discover_port or Discover.DISCOVER_PORT))
|
|
s.listen()
|
|
|
|
while True:
|
|
try:
|
|
conn, _ = s.accept()
|
|
except OSError as e:
|
|
logger.error(f"Accept failed: {e}")
|
|
time.sleep(1)
|
|
continue
|
|
|
|
with conn:
|
|
try:
|
|
conn.settimeout(Discover.CLIENT_TIMEOUT)
|
|
if conn.recv(len(Discover.MAGIC)) == Discover.MAGIC:
|
|
conn.sendall(struct.pack("!H", real_port))
|
|
except OSError as e:
|
|
logger.debug(f"Discovery request failed: {e}")
|
|
|
|
@staticmethod
|
|
def request(target_host, timeout=3, discover_port=None):
|
|
with socket.create_connection((target_host, discover_port or Discover.DISCOVER_PORT), timeout=timeout) as s:
|
|
s.sendall(Discover.MAGIC)
|
|
|
|
data = b""
|
|
while len(data) < 2:
|
|
chunk = s.recv(2 - len(data))
|
|
if not chunk:
|
|
raise ConnectionError("Discovery response was cut short")
|
|
data += chunk
|
|
|
|
return struct.unpack("!H", data)[0]
|