#!/usr/bin/env pmpython
#
# Copyright (C) 2017 Marko Myllynen <myllynen@redhat.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#

try:
    import ConfigParser
except ImportError:
    import configparser as ConfigParser

import os
import pwd
import sys
import socket

from pcp.pmapi import pmContext as PCP

conffile = PCP.pmGetConfig('PCP_PMDAS_DIR')
conffile += '/haproxy/haproxy.conf'

user = 'root'
skt = '/var/lib/haproxy/stats'
url = ''
auth = None
pasw = None

# Python < 3.2 compat
if sys.version_info[0] >= 3 and sys.version_info[1] >= 2:
    config = ConfigParser.ConfigParser()
else:
    config = ConfigParser.SafeConfigParser()
config.read(conffile)
if config.has_section('pmda'):
    for opt in config.options('pmda'):
        if opt == 'user':
            user = config.get('pmda', opt)
        elif opt == 'socket':
            skt = config.get('pmda', opt)
        elif opt == 'url':
            url = config.get('pmda', opt)
        elif opt == 'auth':
            auth = config.get('pmda', opt)
        elif opt == 'pass':
            pasw = config.get('pmda', opt)
        else:
            sys.stderr.write("Invalid directive '%s' in %s.\n" % (opt, conffile))
            sys.exit(1)

if len(sys.argv) > 1 and (sys.argv[1] == '-c' or sys.argv[1] == '--config'):
    sys.stdout.write("user=%s\nsocket=%s\nurl=%s\n" % (user, skt, url.replace(";", "\\;").replace("&", "\\&")))
    sys.exit(0)

try:
    uid = pwd.getpwnam(user).pw_uid
    os.setuid(uid)
except:
    sys.stderr.write("Failed to switch as user %s, try sudo perhaps?\n" % user)
    sys.exit(1)

def setup_urllib(url, auth, pasw):
    """ Setup urllib """
    try:
        import urllib.request as request
    except:
        import urllib2 as request
    passman = request.HTTPPasswordMgrWithDefaultRealm()
    passman.add_password(None, url, auth, pasw)
    authhandler = request.HTTPBasicAuthHandler(passman)
    opener = request.build_opener(authhandler)
    request.install_opener(opener)
    return request

try:
    if not url:
        conn = skt
        stats = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        stats.connect(skt)
        stats.send("show info\n")
        data = stats.recv(4096)
        stats.close()
    else:
        conn = url if not auth else auth + ":<pw>@" + url
        req = setup_urllib(url, auth, pasw).urlopen(url)
        req.read()
        req.close()
except Exception as e:
    sys.stdout.write("Connection as %s to %s failed: %s\n" % (user, conn, e))
    sys.exit(1)

sys.stdout.write("Connection as %s to %s ok.\n" % (user, conn))