Updating known_hosts file periodically for autoadd policy

pull/12/head
Sheng 2018-03-16 12:25:20 +08:00
parent dc0b48ce9a
commit 82a6dfdabd
1 changed files with 24 additions and 15 deletions

39
main.py
View File

@ -6,6 +6,7 @@ import traceback
import uuid import uuid
import weakref import weakref
import paramiko import paramiko
import tornado.ioloop
import tornado.web import tornado.web
import tornado.websocket import tornado.websocket
from tornado.ioloop import IOLoop from tornado.ioloop import IOLoop
@ -19,6 +20,7 @@ define('port', default=8888, help='listen port', type=int)
define('debug', default=False, help='debug mode', type=bool) define('debug', default=False, help='debug mode', type=bool)
define('policy', default='reject', define('policy', default='reject',
help='missing host key policy, reject|autoadd|warning') help='missing host key policy, reject|autoadd|warning')
define('period', default=10, help='seconds for PeriodicCallback', type=int)
BUF_SIZE = 1024 BUF_SIZE = 1024
@ -197,10 +199,8 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
def ssh_connect(self): def ssh_connect(self):
ssh = paramiko.SSHClient() ssh = paramiko.SSHClient()
if isinstance(self.settings['policy'], paramiko.client.AutoAddPolicy): ssh._system_host_keys = self.settings['system_host_keys']
ssh.load_host_keys(self.settings['host_file']) ssh._host_keys = self.settings['host_keys']
else:
ssh._host_keys = self.settings['host_keys']
ssh.set_missing_host_key_policy(self.settings['policy']) ssh.set_missing_host_key_policy(self.settings['policy'])
args = self.get_args() args = self.get_args()
dst_addr = (args[0], args[1]) dst_addr = (args[0], args[1])
@ -290,12 +290,12 @@ def get_host_keys(path):
return paramiko.hostkeys.HostKeys() return paramiko.hostkeys.HostKeys()
def create_host_file(host_file): def save_host_keys(host_keys, filename):
host_keys = get_host_keys(host_file) length = len(host_keys)
if not host_keys: if length != host_keys._last_len:
host_keys = get_host_keys(os.path.expanduser('~/.ssh/known_hosts')) logging.info('Updating {}'.format(filename))
host_keys.save(host_file) host_keys.save(filename)
return host_keys host_keys._last_len = length
def get_policy_class(policy): def get_policy_class(policy):
@ -316,12 +316,21 @@ def get_policy_class(policy):
def main(): def main():
parse_command_line() parse_command_line()
base_dir = os.path.dirname(__file__) base_dir = os.path.dirname(__file__)
host_file = os.path.join(base_dir, 'known_hosts') filename = os.path.join(base_dir, 'known_hosts')
host_keys = create_host_file(host_file) host_keys = get_host_keys(filename)
system_host_keys = get_host_keys(os.path.expanduser('~/.ssh/known_hosts'))
policy_class = get_policy_class(options.policy) policy_class = get_policy_class(options.policy)
if not host_keys and policy_class is paramiko.client.RejectPolicy: if policy_class is paramiko.client.AutoAddPolicy:
raise ValueError('Empty known_hosts with reject policy?') host_keys.save(filename) # for permssion test
host_keys._last_len = len(host_keys)
tornado.ioloop.PeriodicCallback(
lambda: save_host_keys(host_keys, filename),
options.period * 1000 # milliseconds
).start()
elif policy_class is paramiko.client.RejectPolicy:
if not host_keys and not system_host_keys:
raise ValueError('Empty known_hosts with reject policy?')
settings = { settings = {
'template_path': os.path.join(base_dir, 'templates'), 'template_path': os.path.join(base_dir, 'templates'),
@ -337,8 +346,8 @@ def main():
settings.update( settings.update(
debug=options.debug, debug=options.debug,
host_file=host_file,
host_keys=host_keys, host_keys=host_keys,
system_host_keys=system_host_keys,
policy=policy_class() policy=policy_class()
) )