102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
|
||
|
|
# Copyright (c) 2024 Julian Müller (ChaoticByte)
|
||
|
|
# License: MIT
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
|
||
|
|
from argparse import ArgumentParser
|
||
|
|
from pathlib import Path
|
||
|
|
from sys import stdout
|
||
|
|
from sys import stderr
|
||
|
|
|
||
|
|
import asyncssh
|
||
|
|
import yaml
|
||
|
|
|
||
|
|
|
||
|
|
config_host = ""
|
||
|
|
config_port = 8022
|
||
|
|
connected_clients = []
|
||
|
|
config_clients = {
|
||
|
|
# username: asyncssh.SSHAuthorizedKeys
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
class SSHServer(asyncssh.SSHServer):
|
||
|
|
def host_based_auth_supported(self): return False
|
||
|
|
def kbdint_auth_supported(self): return False
|
||
|
|
def password_auth_supported(self): return False
|
||
|
|
def public_key_auth_supported(self): return True
|
||
|
|
def begin_auth(self, username: str) -> bool: return True # we wanna handle auth
|
||
|
|
|
||
|
|
def validate_public_key(self, username: str, key: asyncssh.SSHKey) -> bool:
|
||
|
|
try:
|
||
|
|
return config_clients[username].validate(key, "", "") is not None
|
||
|
|
except:
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
def broadcast(msg: str):
|
||
|
|
assert type(msg) == str
|
||
|
|
for c in connected_clients:
|
||
|
|
c.stdout.write(msg)
|
||
|
|
|
||
|
|
async def handle_connection(process: asyncssh.SSHServerProcess):
|
||
|
|
connected_clients.append(process)
|
||
|
|
username = process.get_extra_info("username")
|
||
|
|
try:
|
||
|
|
connected_msg = f"[connected] {username}\n"
|
||
|
|
stdout.write(connected_msg)
|
||
|
|
broadcast(connected_msg)
|
||
|
|
while True:
|
||
|
|
try:
|
||
|
|
async for line in process.stdin:
|
||
|
|
if line == "": raise asyncssh.BreakReceived(0)
|
||
|
|
line = line.strip('\n\r')
|
||
|
|
msg = f"{username}: {line}\n"
|
||
|
|
stdout.write(msg)
|
||
|
|
broadcast(msg)
|
||
|
|
except asyncssh.TerminalSizeChanged:
|
||
|
|
continue
|
||
|
|
finally:
|
||
|
|
break
|
||
|
|
except asyncssh.BreakReceived:
|
||
|
|
pass
|
||
|
|
except Exception as e:
|
||
|
|
stderr.write(f"An error occured: {type(e).__name__} {e}\n")
|
||
|
|
stderr.flush()
|
||
|
|
finally:
|
||
|
|
process.exit(0)
|
||
|
|
connected_clients.remove(process)
|
||
|
|
disconnected_msg = f"[disconnected] {username}\n"
|
||
|
|
stdout.write(disconnected_msg)
|
||
|
|
broadcast(disconnected_msg)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
# commandline arguments
|
||
|
|
argp = ArgumentParser()
|
||
|
|
argp.add_argument("config", type=Path, help="The path to the config file")
|
||
|
|
args = argp.parse_args()
|
||
|
|
# read config
|
||
|
|
config = yaml.safe_load(args.config.read_text())
|
||
|
|
config_host = str(config["host"])
|
||
|
|
config_port = int(config["port"])
|
||
|
|
config_private_key = asyncssh.import_private_key(config["private_key"])
|
||
|
|
server_public_key = config_private_key.export_public_key("openssh").decode().strip("\n\r")
|
||
|
|
stderr.write(f"Server public key is \"{server_public_key}\"\n")
|
||
|
|
stderr.flush()
|
||
|
|
for c in config["clients"]:
|
||
|
|
config_clients[str(c)] = asyncssh.import_authorized_keys(str(config["clients"][c]))
|
||
|
|
# start server
|
||
|
|
loop = asyncio.get_event_loop()
|
||
|
|
loop.run_until_complete(
|
||
|
|
asyncssh.create_server(
|
||
|
|
SSHServer,
|
||
|
|
config_host,
|
||
|
|
config_port,
|
||
|
|
server_host_keys=[config_private_key],
|
||
|
|
process_factory=handle_connection
|
||
|
|
))
|
||
|
|
loop.run_forever()
|