test_mediaworker.py

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

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


try:
    import mediaserver
except ImportError:
    print('MediaServer is not installed, skipping test.')
    sys.exit(2)
else:
    print('MediaServer version: %s.' % mediaserver.__version__)


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 to MediaWorker:\n%s' % cmd)
    try:
        subprocess.check_output(cmd, shell=True, timeout=2)
        print('%sLogged in successfully in "%s".%s' % (GREEN, ip, DEF))
    except subprocess.CalledProcessError:
        print('%sFailed to login using SSH, run "ssh-copy-id %s".%s' % (RED, ip, DEF))
        return False
    return True


def run_tests(ip):
    print('Updating envsetup tests on MediaWorker.')
    cmd = 'ssh -t %s /root/envsetup/update_envsetup.py' % ip
    os.system(cmd)
    print('Running envsetup tests on MediaWorker.')
    cmd = 'ssh -t %s /root/envsetup/tester.py' % ip
    status = os.system(cmd)
    if status == 0:
        subprocess.check_output(cmd, shell=True, timeout=60)
        print('%sAll tests completed on MediaWorker "%s".' % (GREEN, ip, DEF))
        return True
    else:
        print('%sapt-get update failed on MediaWorker "%s".%s' % (RED, ip, DEF))
        return False


def check_celerity_connectivity(ip):
    print('Getting celerity server url.')
    cmd = 'ssh -t %s cat /etc/celerity/config.py' % ip
    print(cmd)
    try:
        d = subprocess.check_output(cmd, shell=True, timeout=5, universal_newlines=True)
    except subprocess.CalledProcessError:
        m = None
    else:
        m = re.search(r'\nSERVER_URL\s*=\s*([:\/\'\"\-\_\.\w]+)', d)
    if not m:
        print('%sFailed to get celerity tasks server url from configuration in MediaWorker "%s".%s' % (RED, ip, DEF))
        return False
    server_url = m.groups()[0].strip('"\' ')
    print('Checking celerity connectivity.')
    cmd = 'ssh -t %s curl -k %s' % (ip, server_url)
    print(cmd)
    try:
        d = subprocess.check_output(cmd, shell=True, timeout=5, universal_newlines=True)
    except subprocess.CalledProcessError:
        d = ''
    if 'Celerity tasks server' in d:
        print('%sSuccessfully reached tasks server from MediaWorker "%s".%s' % (GREEN, ip, DEF))
        return True
    print('%sFailed to reach tasks server from MediaWorker "%s".%s' % (RED, ip, DEF))
    return False


def check_mediaworker_in_whitelist(ip):
    nginx_vhosts_path = '/etc/nginx/sites-enabled'
    vhosts = os.listdir(nginx_vhosts_path)
    for v in vhosts:
        if v.endswith('.conf'):
            vhost_path = os.path.join(nginx_vhosts_path, v)
            with open(vhost_path, 'r') as f:
                d = f.read()
                if 'msuser_whitelist' in d:
                    if ip not in d:
                        print('%sMediaWorker ip %s is not in %s whitelist%s' % (RED, ip, v, DEF))
                        return False
    return True

all_ok = True
tested = False
mediaserver_ip = conf.get('NETWORK_IP')
worker_ips = conf.get('CELERITY_WORKER_IP')
for worker_ip in worker_ips.split(','):
    worker_ip = worker_ip.strip()
    if worker_ip and worker_ip != '127.0.1.1' and worker_ip != mediaserver_ip:
        tested = True
        if not check_ssh(worker_ip):
            all_ok = False
        else:
            if not check_celerity_connectivity(worker_ip):
                all_ok = False
            # if not run_tests(worker_ip):
            #     all_ok = False
            if not check_mediaworker_in_whitelist(worker_ip):
                all_ok = False
if not tested:
    print('Celerity IP not set or running locally, skipping test.')
    sys.exit(2)

sys.exit(int(not all_ok))