#!/skynet/python/bin/python

"""Collects information about currently active VLANs on hosts.

Required by check-vlan-configuration script.
"""

from __future__ import print_function
from __future__ import unicode_literals

import os
import json
import re
import subprocess

from api.cqueue import Client
from library.sky.hosts import resolveHosts


def get_active_vlans():
    ifconfig = subprocess.Popen(["ifconfig"], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True)
    stdout, stderr = ifconfig.communicate()

    vlans = sorted({int(vlan.group(1)) for vlan in re.finditer(r"^vlan(\d+)", stdout, re.MULTILINE)})
    return (vlans, stderr)

get_active_vlans.osUser = "loadbase"


def iter_chunks(data, size):
    for i in xrange(0, len(data), size):
        yield data[i:i+size]


def main():
    vlan_map = {}

    with Client() as client:
        with open(os.path.expanduser("~/.ssh/id_loadbase"), "rb") as key_file:
            key = client.signer.loadKeys(key_file, "id_loadbase").next()

        client.signer.addKey(key)

        hosts = list(resolveHosts(["i+ALL_SEARCH"])[0])

        for hosts_chunk in iter_chunks(hosts, 1000):
            with client.run(hosts_chunk, get_active_vlans) as session:
                for host, result, error in session.wait(timeout=30):
                    if error is None:
                        result, error = result

                    if error:
                        print(host + ":", error)
                    else:
                        vlan_map[host] = result

            print(len(vlan_map))

    with open("vlan_map.json", "w") as out:
        json.dump(vlan_map, out)


if __name__ == "__main__":
    main()