test_mediaworker.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
Criticality: High
Checks that MediaWorker can be reached using SSH
'''
import os
import imp
import sys
import subprocess

GREEN = '\033[92m'
RED = '\033[91m'
DEF = '\033[0m'

os.chdir(os.path.dirname(__file__))
if not os.path.isfile('../utils.py'):
    print('conf.sh not found')
    sys.exit(1)

es_utils = imp.load_source('es_utils', '../utils.py')
conf = es_utils.load_conf()

def check_ssh(ip):
    cmd = 'ssh -o StrictHostKeyChecking=no -o PasswordAuthentication=no %s ls /tmp' % ip
    print('Connecting: %s' % cmd)
    try:
        subprocess.check_output(cmd, shell=True, timeout=2)
        print('%sLogged in successfully%s' % (GREEN, DEF))
    except subprocess.CalledProcessError:
        print('%sFailed to login using SSH, run ssh-copy-id %s %s' % (RED, ip, DEF))
        return False
    return True

def check_apt(ip):
    cmd = 'ssh %s apt-get update' % ip
    try:
        subprocess.check_output(cmd, shell=True, timeout=10)
        print('%sapt-get update success on MediaWorker%s' % (GREEN, DEF))
    except subprocess.CalledProcessError:
        print('%apt-get update failed on MediaWorker%s' % (RED, DEF))
        return False

all_ok = True

worker_ip = conf.get('CELERITY_WORKER_IP')
if worker_ip != '127.0.1.1':
    if not check_ssh(worker_ip):
        all_ok = False
    if not check_apt(worker_ip):
        all_ok = False
sys.exit(int(not all_ok))