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 root@%s ls /tmp' % ip
    print('Connecting to MediaWorker:\n%s' % cmd)
    try:
        subprocess.check_output(cmd, shell=True, timeout=5)
        print('%sLogged in successfully in "%s".%s' % (GREEN, ip, DEF))
    except subprocess.CalledProcessError:
        # do some diagnostics
        try:
            cmd_port = 'nc -z -w2 {} 22'.format(ip)
            subprocess.check_output(cmd_port, shell=True, timeout=5)
        except subprocess.CalledProcessError:
            cmd_ping = 'ping -c2 -w4 {}'.format(ip)
            subprocess.check_output(cmd_ping, shell=True, timeout=5)
        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 root@%s /root/envsetup/update_envsetup.py' % ip
    os.system(cmd)
    print('Running envsetup tests on MediaWorker.')
    cmd = 'ssh -t root@%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 root@%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 root@%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_celerity_versions(ip):
    print('Checking tha celerity server and worker uses the same version.')
    try:
        ms_out = subprocess.check_output('dpkg -s celerity-utils | grep "^Version:"', shell=True, timeout=10, universal_newlines=True)
        mw_out = subprocess.check_output('ssh -t root@%s dpkg -s celerity-utils | grep "^Version:"' % ip, shell=True, timeout=10, universal_newlines=True)
    except subprocess.CalledProcessError as e:
        print('%sFailed to check celerity version in MediaWorker "%s":\n%s%s' % (RED, ip, e, DEF))
        return False
    ms_out = (ms_out[len('Version:'):] if ms_out.startswith('Version:') else ms_out).strip()
    mw_out = (mw_out[len('Version:'):] if mw_out.startswith('Version:') else mw_out).strip()
    if ms_out != mw_out:
        print('%sThe celerity version in MediaWorker "%s" is not the same as in MediaServer.%s\nMediaServer version: \tn%s\nMediaWorker version: \t%s' % (RED, ip, DEF, ms_out, mw_out))
        return False
    print('%sThe celerity version in MediaWorker "%s" is the same as in MediaServer.%s\nCurrent celerity version is: %s.' % (GREEN, ip, DEF, ms_out))
    return True


def check_mediaworker_in_whitelist(ip):
    # this check is not usefull anymore because the worker uses the
    # API to get links to resources with valid secure link token
    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 not worker_ip.startswith('127.0.') 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 check_celerity_versions(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))