| #!/usr/bin/python3 -u |
| |
| import argparse |
| import os |
| import re |
| import requests |
| import sys |
| import socket |
| import struct |
| import time |
| |
| from multiprocessing import Process |
| from multiprocessing import Event |
| from http.server import HTTPServer |
| from http.server import BaseHTTPRequestHandler |
| |
| import netifaces |
| |
| |
| USERNAME = 'mendel' |
| GROUP = 'mendel' |
| SSH_BASE_PATH = '/home/mendel/.ssh' |
| AUTHORIZED_KEYS_PATH = os.path.join(SSH_BASE_PATH, 'authorized_keys') |
| SERVER_PORT = 41337 |
| MDT_KEY_REGEX = re.compile(r'^ssh-rsa .* mdt$') |
| MDT_EXIT_REGEX = re.compile(r'^exit$') |
| |
| |
| def get_iface_address(iface): |
| """ Retreive ip address from a given interface. """ |
| addresses = netifaces.ifaddresses(iface) |
| inet4_addresses = addresses[socket.AF_INET] |
| ip_address = inet4_addresses[0]['addr'] |
| return ip_address |
| |
| |
| def key_exists(): |
| """ Check for existence of key. """ |
| if os.path.exists(AUTHORIZED_KEYS_PATH): |
| print('authorized_keys file already exists.\n', flush=True) |
| with open(AUTHORIZED_KEYS_PATH) as fp: |
| for line in fp.readlines(): |
| if re.match(MDT_KEY_REGEX, line): |
| print('authorized_keys file contains MDT key already.\n', flush=True) |
| return True |
| print('authorized_keys file does not contain an MDT key.\n', flush=True) |
| return False |
| |
| |
| class KeyPushHandler(BaseHTTPRequestHandler): |
| """ Handle keys sent from client and install. """ |
| key_handled = Event() |
| def do_PUT(self): |
| """ Verify and install key. """ |
| if not KeyPushHandler.key_handled.is_set(): |
| content_length = int(self.headers.get('Content-Length', 0)) |
| public_key = self.rfile.read(content_length) |
| if not re.match(MDT_KEY_REGEX, public_key.decode("utf-8").strip()): |
| if re.match(MDT_EXIT_REGEX, public_key.decode("utf-8").strip()): |
| KeyPushHandler.key_handled.set() |
| print("Received request to exit.", flush=True) |
| self.send_response(200, "Ok") |
| self.end_headers() |
| return |
| KeyPushHandler.key_handled.clear() |
| print("Received invalid key.", flush=True) |
| self.send_response(400, "Invalid key") |
| self.end_headers() |
| return |
| else: |
| KeyPushHandler.key_handled.set() |
| print("Received valid key.", flush=True) |
| |
| if not os.path.exists(SSH_BASE_PATH): |
| os.mkdir(SSH_BASE_PATH, 0o700) |
| |
| with open(AUTHORIZED_KEYS_PATH, 'ab') as fp: |
| fp.write(public_key) |
| |
| sys.stdout.write('authorized_keys file written -- exiting.\n') |
| sys.stdout.flush() |
| |
| self.send_response(200, "Ok") |
| self.end_headers() |
| |
| |
| class HTTPServerProcess(Process): |
| """ Process for starting a HttpServer """ |
| def __init__(self, bind_address, *args, **kwargs): |
| super(HTTPServerProcess, self).__init__(*args, **kwargs) |
| self.bind_address = bind_address |
| |
| def run(self): |
| """ Bind and start server on address.""" |
| server_address = (self.bind_address, SERVER_PORT) |
| httpd = HTTPServer(server_address, KeyPushHandler) |
| print('Waiting for incoming PUT on {0}:{1}\n'.format(self.bind_address, |
| SERVER_PORT), flush=True) |
| while not KeyPushHandler.key_handled.is_set(): |
| httpd.handle_request() |
| |
| def send_exit_command(self): |
| """ Send command to server to exit. """ |
| print("Request exit for {0}".format(self.bind_address, SERVER_PORT)) |
| try: |
| requests.put("http://{0}:{1}".format(self.bind_address, SERVER_PORT), data="exit") |
| except requests.exceptions.ConnectionError: |
| pass |
| |
| |
| def main(): |
| parser = argparse.ArgumentParser(description="MDT Keymaster Server.") |
| parser.add_argument("--interface", help="Listen on selected interfaces.", |
| action="append", default=[]) |
| args = parser.parse_args() |
| |
| interfaces = list(set(args.interface)) |
| |
| # If key is installed exit. |
| if key_exists(): |
| print('exiting.\n', flush=True) |
| sys.exit(0) |
| |
| ifaces = args.interface |
| bind_addresses = list() |
| delay = 1 |
| while len(bind_addresses) != len(ifaces): |
| for iface in ifaces: |
| try: |
| bind_addresses.append(get_iface_address(iface)) |
| except Exception as e: |
| print('Unable to determine {0} IPv4 address: {1}.\n'.format(iface, e), flush=True) |
| print('Waiting {0} seconds before retrying.'.format(delay), flush=True) |
| time.sleep(delay) |
| delay = min(10, delay * 2) |
| |
| server_processes = [] |
| for bind_address in bind_addresses: |
| server_processes.append((bind_address, |
| HTTPServerProcess(bind_address))) |
| |
| for address, process in server_processes: |
| process.start() |
| |
| while not KeyPushHandler.key_handled.is_set(): |
| time.sleep(1.0) |
| |
| for address, process in server_processes: |
| process.send_exit_command() |
| process.join() |
| |
| print('Exiting.\n', flush=True) |
| sys.exit(0) |
| |
| |
| if __name__ == '__main__': |
| main() |