343 lines
12 KiB
Python
343 lines
12 KiB
Python
"""collabore tunnel SSH server"""
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import random
|
|
import string
|
|
import sys
|
|
import time
|
|
from asyncio import AbstractEventLoop
|
|
from collections import deque
|
|
from os import path
|
|
from types import FrameType
|
|
from typing import AnyStr, Optional, Tuple
|
|
from _asyncio import Task
|
|
|
|
import asyncssh
|
|
from asyncssh import SSHKey, SSHServerConnection
|
|
from asyncssh.channel import (
|
|
SSHUNIXChannel,
|
|
SSHUNIXSession,
|
|
SSHUNIXSessionFactory,
|
|
)
|
|
from asyncssh.listener import create_unix_forward_listener
|
|
from asyncssh.misc import MaybeAwait
|
|
from loguru import logger
|
|
from loguru._handler import Handler
|
|
|
|
|
|
unix_sockets_dir: str = os.getenv("UNIX_SOCKETS_DIRECTORY", "/tmp/collabore-tunnel")
|
|
server_hostname: str = os.getenv("SERVER_HOSTNAME", "tnl.clb.re")
|
|
config_dir: str = os.getenv("CONFIG_DIRECTORY", ".")
|
|
welcome_banner_file: str = os.getenv("WELCOME_BANNER_FILE", "./welcome_banner.txt")
|
|
rate_limit_count: int = int(os.getenv("RATE_LIMIT_COUNT", "5"))
|
|
rate_limit_interval: int = int(os.getenv("RATE_LIMIT_INTERVAL", "60"))
|
|
max_connections_per_ip: int = int(os.getenv("MAX_CONNECTIONS_PER_IP", "5"))
|
|
timeout: int = int(os.getenv("TIMEOUT", "120"))
|
|
ssh_server_host: str = os.getenv("SSH_SERVER_HOST", "0.0.0.0")
|
|
ssh_server_port: int = int(os.getenv("SSH_SERVER_PORT", "22"))
|
|
log_level: str = os.getenv("LOG_LEVEL", "INFO")
|
|
log_depth: int = int(os.getenv("LOG_DEPTH", "2"))
|
|
|
|
|
|
def read_welcome_banner() -> str:
|
|
"""Read the welcome banner from a file"""
|
|
if not os.path.exists(welcome_banner_file):
|
|
return welcome_banner
|
|
with open(welcome_banner_file, "r", encoding="UTF-8") as file:
|
|
return file.read()
|
|
|
|
|
|
welcome_banner: str = read_welcome_banner()
|
|
|
|
|
|
class RateLimiter:
|
|
"""Rate limiter handling class"""
|
|
|
|
def __init__(self, max_requests: int, interval: int):
|
|
"""Init class"""
|
|
self.max_requests: int = max_requests
|
|
self.interval: int = interval
|
|
self.timestamps: deque = deque()
|
|
|
|
def is_rate_limited(self) -> bool:
|
|
"""Check if rate limited"""
|
|
now: float = time.time()
|
|
while self.timestamps and self.timestamps[0] < now - self.interval:
|
|
self.timestamps.popleft()
|
|
if len(self.timestamps) >= self.max_requests:
|
|
return True
|
|
self.timestamps.append(now)
|
|
return False
|
|
|
|
|
|
class ConcurrentConnections:
|
|
"""Concurrent connection handling class"""
|
|
|
|
def __init__(self):
|
|
"""Init class"""
|
|
self.ip_connections: dict = {}
|
|
|
|
def increment(self, ip_addr: str) -> None:
|
|
"""Increment the number of concurrent connections for an IP"""
|
|
if ip_addr not in self.ip_connections:
|
|
self.ip_connections[ip_addr] = 1
|
|
else:
|
|
self.ip_connections[ip_addr] += 1
|
|
|
|
def decrement(self, ip_addr: str) -> None:
|
|
"""Decrement the number of concurrent connections for an IP"""
|
|
self.ip_connections[ip_addr] -= 1
|
|
|
|
def get(self, ip_addr: str) -> int:
|
|
"""Get the number of concurent connection for an IP"""
|
|
return self.ip_connections.get(ip_addr, 0)
|
|
|
|
|
|
ip_address_connections = ConcurrentConnections()
|
|
|
|
|
|
def check_concurrent_connections(ip_addr: str) -> bool:
|
|
"""Checking for concurrent connections"""
|
|
return ip_address_connections.get(ip_addr) >= max_connections_per_ip
|
|
|
|
|
|
class SSHServer(asyncssh.SSHServer):
|
|
"""SSH server protocol handler class"""
|
|
|
|
rate_limiters: dict = {}
|
|
|
|
def __init__(self):
|
|
"""Init class"""
|
|
self.conn: SSHServerConnection
|
|
self.socket_path: str
|
|
self.ip_addr: str
|
|
|
|
def check_rate_limit(self, ip_addr: str) -> bool:
|
|
"""Check if rate limited"""
|
|
if ip_addr not in self.rate_limiters:
|
|
self.rate_limiters[ip_addr] = RateLimiter(
|
|
rate_limit_count, rate_limit_interval
|
|
)
|
|
return self.rate_limiters[ip_addr].is_rate_limited()
|
|
|
|
def connection_made(self, conn: SSHServerConnection) -> None:
|
|
"""Called when a connection is made"""
|
|
self.conn = conn
|
|
self.ip_addr, _ = conn.get_extra_info("peername")
|
|
|
|
if self.check_rate_limit(self.ip_addr):
|
|
conn.set_extra_info(rate_limited=True)
|
|
|
|
if check_concurrent_connections(self.ip_addr):
|
|
conn.set_extra_info(connection_limited=True)
|
|
|
|
ip_address_connections.increment(self.ip_addr)
|
|
|
|
def connection_lost(self, exc: Optional[Exception]) -> None:
|
|
"""Called when a connection is lost or closed"""
|
|
if exc:
|
|
logging.info("The connection has been terminated: %s", str(exc))
|
|
try:
|
|
os.remove(self.socket_path)
|
|
except AttributeError:
|
|
pass
|
|
ip_address_connections.decrement(self.ip_addr)
|
|
|
|
def begin_auth(self, username: str) -> MaybeAwait[bool]:
|
|
"""Authentication has been requested by the client"""
|
|
return False
|
|
|
|
def password_auth_supported(self) -> bool:
|
|
"""Return whether or not password authentication is supported"""
|
|
return True
|
|
|
|
def generate_socket_path(self) -> str:
|
|
"""Return the path of a socket whose name has been randomly generated"""
|
|
socket_name = get_random_slug(16)
|
|
self.socket_path = os.path.join(unix_sockets_dir, f"{socket_name}.sock")
|
|
self.conn.set_extra_info(socket_name=socket_name)
|
|
return self.socket_path
|
|
|
|
def unix_server_requested(self, listen_path: str):
|
|
"""Handle a request to listen on a UNIX domain socket"""
|
|
rewrite_path: str = self.generate_socket_path()
|
|
|
|
async def tunnel_connection(
|
|
session_factory: SSHUNIXSessionFactory[AnyStr],
|
|
) -> Tuple[SSHUNIXChannel[AnyStr], SSHUNIXSession[AnyStr]]:
|
|
return await self.conn.create_unix_connection(session_factory, listen_path)
|
|
|
|
try:
|
|
return create_unix_forward_listener(
|
|
self.conn, asyncio.get_event_loop(), tunnel_connection, rewrite_path
|
|
)
|
|
except OSError as create_unix_forward_listener_exception:
|
|
logging.error(
|
|
"An error occurred while creating the forward listener: %s",
|
|
str(create_unix_forward_listener_exception),
|
|
)
|
|
|
|
|
|
async def handle_ssh_client(process) -> None:
|
|
"""Function called every time a client connects to the SSH server"""
|
|
socket_name: str = process.get_extra_info("socket_name")
|
|
rate_limited: bool = process.get_extra_info("rate_limited")
|
|
connection_limited: bool = process.get_extra_info("connection_limited")
|
|
response: str = ""
|
|
|
|
async def process_timeout(process):
|
|
"""Function to terminate the connection automatically
|
|
after a specific period of time (in minutes)"""
|
|
await asyncio.sleep(timeout * 60)
|
|
response = (
|
|
f"Timeout: you were automatically ejected after {timeout} minutes of use.\n"
|
|
)
|
|
process.stdout.write(response + "\n")
|
|
process.logger.info(
|
|
f"The user was automatically ejected after {timeout} minutes of use"
|
|
)
|
|
process.close()
|
|
|
|
if not rate_limited:
|
|
if not connection_limited:
|
|
if not socket_name:
|
|
response = "Usage: ssh -R /:host:port ssh.tunnel.collabore.fr\n"
|
|
process.stdout.write(response + "\n")
|
|
process.logger.info(
|
|
"The user was ejected because they did not connect in port forwarding mode"
|
|
)
|
|
process.exit(1)
|
|
return
|
|
no_tls: str = f"{socket_name}.{server_hostname}"
|
|
tls: str = f"https://{socket_name}.{server_hostname}"
|
|
response = f"{welcome_banner}\nYour local service has been exposed to the public\n\
|
|
Internet address: {no_tls}\nTLS termination: {tls}\n"
|
|
process.stdout.write(response + "\n")
|
|
process.logger.info(f"Exposed on {no_tls}")
|
|
read_task: Task = asyncio.create_task(process.stdin.read())
|
|
timeout_task: Task = asyncio.create_task(process_timeout(process))
|
|
done, pending = await asyncio.wait(
|
|
[read_task, timeout_task], return_when=asyncio.FIRST_COMPLETED
|
|
)
|
|
for task in done:
|
|
try:
|
|
await task
|
|
except asyncssh.BreakReceived:
|
|
pass
|
|
for task in pending:
|
|
task.cancel()
|
|
|
|
process.exit(0)
|
|
else:
|
|
response = (
|
|
"Per-IP connection limit: too many connections running over this IP.\n"
|
|
)
|
|
process.stdout.write(response + "\n")
|
|
process.logger.warning("Rejected connection due to per-IP connection limit")
|
|
process.exit(1)
|
|
return
|
|
else:
|
|
response = "Rate limited: please try later.\n"
|
|
process.stdout.write(response + "\n")
|
|
process.logger.warning("Rejected connection due to rate limit")
|
|
process.exit(1)
|
|
return
|
|
|
|
|
|
async def start_ssh_server() -> None:
|
|
"""Start the SSH server"""
|
|
ssh_key_file: str = path.join(config_dir, "id_rsa_host")
|
|
await asyncssh.create_server(
|
|
SSHServer,
|
|
host=ssh_server_host,
|
|
port=ssh_server_port,
|
|
server_host_keys=[ssh_key_file],
|
|
process_factory=handle_ssh_client,
|
|
agent_forwarding=False,
|
|
allow_scp=False,
|
|
keepalive_interval=30,
|
|
)
|
|
logging.info("SSH server started successfully.")
|
|
|
|
|
|
def check_unix_sockets_dir() -> None:
|
|
"""If the directory for UNIX sockets does not exist, it is created"""
|
|
if not path.exists(unix_sockets_dir):
|
|
os.mkdir(unix_sockets_dir)
|
|
logging.warning(
|
|
"The %s folder does not exist, it has been created.", unix_sockets_dir
|
|
)
|
|
else:
|
|
logging.info("The %s folder exist.", unix_sockets_dir)
|
|
|
|
|
|
def generate_ssh_key() -> None:
|
|
"""If the SSH key of the server does not exist, it is generated"""
|
|
ssh_host_key: str = path.join(config_dir, "id_rsa_host")
|
|
logging.info("Loading the SSH key")
|
|
if not path.exists(ssh_host_key):
|
|
logging.warning(
|
|
"The SSH key for the host was not found, generation in progress..."
|
|
)
|
|
key: SSHKey = asyncssh.generate_private_key("ssh-rsa")
|
|
private_key: bytes = key.export_private_key()
|
|
with open(ssh_host_key, "wb") as ssh_host_key_data:
|
|
ssh_host_key_data.write(private_key)
|
|
logging.info("The key was successfully created!")
|
|
else:
|
|
logging.info("SSH key has been found")
|
|
|
|
|
|
class InterceptHandler(logging.Handler):
|
|
"""Intercept logging call"""
|
|
|
|
def emit(self, record):
|
|
"""Find caller from where originated the logged message"""
|
|
frame: FrameType = logging.currentframe()
|
|
depth: int = log_depth
|
|
while frame.f_code.co_filename == logging.__file__:
|
|
frame = frame.f_back
|
|
depth += 1
|
|
logger.opt(exception=record.exc_info).log(log_level, record.getMessage())
|
|
|
|
|
|
def init_logging():
|
|
"""Init logging with a custom handler"""
|
|
logging.root.handlers: Handler = [InterceptHandler()]
|
|
logging.root.setLevel(log_level)
|
|
fmt = "<green>[{time}]</green> <level>[{level}]</level> - <level>{message}</level>"
|
|
logger.configure(handlers=[{"sink": sys.stdout, "serialize": False, "format": fmt}])
|
|
|
|
|
|
def get_random_slug(length) -> str:
|
|
"""Function that generates a random string of a defined size"""
|
|
chars: str = string.ascii_lowercase + string.digits
|
|
return "".join(random.choices(chars, k=length))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
init_logging()
|
|
logging.info("Starting collabore tunnel SSH server...")
|
|
os.umask(0o000)
|
|
generate_ssh_key()
|
|
logging.info("Checking for the existence of a folder for UNIX sockets...")
|
|
check_unix_sockets_dir()
|
|
loop: AbstractEventLoop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(start_ssh_server())
|
|
except KeyboardInterrupt:
|
|
pass
|
|
except (OSError, asyncssh.Error) as ssh_server_startup_exception:
|
|
logging.critical(
|
|
"An error occurred while starting the SSH server: %s",
|
|
str(ssh_server_startup_exception),
|
|
)
|
|
sys.exit()
|
|
try:
|
|
loop.run_forever()
|
|
except KeyboardInterrupt:
|
|
sys.exit()
|