#!/usr/bin/env python

from __future__ import with_statement
import argparse
import functools
import logging
import os
import os.path
import re
import sys
import urllib2
import urlparse
import pexpect

logger = logging.getLogger(__name__)
logger.addHandler(logging.StreamHandler(sys.stderr))

USE_PROXY = True
HTTP_PROXY = "http://squid-6dc7f4.sfo01.justin.tv:9797/" # for urllib2
SOCKS_PROXY = "prxy.internal.justin.tv:1080" # for scp

DESCRIPTION = """Checks that production SWFs on Twitch are whitelisted by L3.
L3 will reject video connections from any SWF not in its whitelist. SWFs
missing from the whitelist indicate a serious Quality of Service issue.
"""

SCP_PROXY = "-o ProxyCommand='nc -x {0} %h %p'".format(SOCKS_PROXY)
SCP_WHITELIST = "scp {0} justintv22@justintv2.ingest.cdn.level3.net:/published/SWF2/dirindex.whitelist {1}"
SCP_PASSWORD = "x7yaInehFA"

WHITELIST_SWF_RE = re.compile(r'^\s*#\s*([^/\.]+\.[a-z0-9]+)\.swf\s*$', re.IGNORECASE)

# We call SWF_URL.format() which is why we have the trailing {0}
SWF_URL = "http://www.twitch.tv{0}"
SWF_VERSION_RE = re.compile(r'\.([a-z0-9]+)\.swf(\?.*)*$', re.IGNORECASE)

def alert_nagios_on_error(func):
    @functools.wraps(func)
    def wrapped(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            print("[FAIL] Unexpected error while validating L3 SWF whitelist|{0}".format(e))
            sys.exit(2)
    return wrapped

@alert_nagios_on_error
def main():
    parser = argparse.ArgumentParser(description=DESCRIPTION)
    parser.add_argument('file_url',
            help="URL of paths file. SWF paths to check (each path separated by newlines)")
    parser.add_argument('-v', dest='verbose', action='store_const', const=True, default=False,
            help='Enable verbose mode')
    args = parser.parse_args()
    logger.setLevel(logging.DEBUG if args.verbose else logging.WARNING)

    paths = paths_to_check(args.file_url)
    if paths is None:
        print("[FAIL] Failed to retrieve L3 SWF paths|SWF Paths url:{0}".format(args.file_url))
        sys.exit(2)

    missing_swfs = []
    non_whitelisted_swfs = []

    if paths:
        whitelist = whitelisted_level3_swfs()
        logger.debug("Found {0} SWFs in L3 Whitelist".format(len(whitelist)))

        urlopener = urllib2.build_opener(SWFRedirectHandler())
        for path in paths:
            version = production_swf_version(urlopener, path)
            if version:
                if version not in whitelist:
                    logger.debug("FAIL [WHITELIST] {0}".format(path))
                    non_whitelisted_swfs.append(path)
                else:
                    logger.debug("OK {0}".format(path))
            else:
                logger.debug("FAIL [ERROR] {0}".format(path))
                missing_swfs.append(path)

    if missing_swfs or non_whitelisted_swfs:
        error_output = []
        error_output.append("[CRITICAL] SWF L3 Whitelist|")
        if missing_swfs:
            error_output.append("SWFs not found on Twitch:")
            for swf in missing_swfs:
                error_output.append(swf)
        if non_whitelisted_swfs:
            error_output.append("SWFs not whitelisted on L3:")
            for swf in non_whitelisted_swfs:
                error_output.append(swf)
        print("\n".join(error_output) + "|")
        sys.exit(2)
    else:
        print("[OK] SWF L3 Whitelist|")
        sys.exit(0)

def whitelisted_level3_swfs():
    temp = "/tmp/l3swfs.whitelist"
    download_command = SCP_WHITELIST.format(SCP_PROXY, temp)
    scp = pexpect.spawn('/bin/bash', ['-c', download_command])
    scp.expect('password:')
    scp.sendline(SCP_PASSWORD)
    scp.expect(pexpect.EOF)

    whitelisted_swfs = set([])
    with open(temp) as whitelist:
        for line in whitelist:
            match = WHITELIST_SWF_RE.match(line)
            if match:
                whitelisted_swfs.add(match.group(1))
    os.remove(temp)

    return whitelisted_swfs

def production_swf_version(urlopener, path):
    url = SWF_URL.format(path)
    try:
        response = urlopener.open(url)
    except (urllib2.URLError, urllib2.HTTPError) as e:
        logger.debug("{0} [{1}]".format(e, url))
        return None

    redirect = response.headers.get('location', None)
    if not redirect:
        return None
    return parse_version_from_url(redirect)

def parse_version_from_url(url):
    match = SWF_VERSION_RE.search(url)
    if not match:
        return None
    parsed = urlparse.urlparse(url)
    return parse_version_from_path(os.path.basename(parsed.path))

def parse_version_from_path(path):
    return os.path.splitext(path)[0]

def paths_to_check(pathsfile_url):
    try:
        response = urllib2.urlopen(pathsfile_url)
    except (urllib2.URLError, urllib2.HTTPError) as e:
        logger.debug("{0} [{1}]".format(e, pathsfile_url))
        return None
    return nonblank_lines(response)

def nonblank_lines(fp):
    paths = [line.strip() for line in fp]
    return [path for path in paths if path]

class SWFRedirectHandler(urllib2.HTTPRedirectHandler):
    def http_error_302(self, req, fp, code, msg, hdrs):
        return fp

if __name__ == '__main__':
    if USE_PROXY:
        # if USE_PROXY is true, we manually set the proxy
        os.environ['http_proxy'] = HTTP_PROXY
    else:
        # if its false, we need to unset SCP_PROXY command
        SCP_PROXY = ''
    main()
