#!/usr/bin/python

# (c) 2012, Michael DeHaan <michael.dehaan@gmail.com>
#
# This file is part of Ansible
#
# Ansible 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 3 of the License, or
# (at your option) any later version.
#
# Ansible 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.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.

import array
import fcntl
import glob
import sys
import os
import platform
import re
import shlex
import socket
import struct
import subprocess
import traceback
import syslog

try:
    import selinux
    HAVE_SELINUX=True
except ImportError:
    HAVE_SELINUX=False

try:
    import json
except ImportError:
    import simplejson as json

class Facts(object):
    """
    This class should only attempt to populate those facts that
    are mostly generic to all systems.  This includes platform facts,
    service facts (eg. ssh keys or selinux), and distribution facts.
    Anything that requires extensive code or may have more than one
    possible implementation to establish facts for a given topic should
    subclass Facts.
    """

    _I386RE = re.compile(r'i[3456]86')
    # For the most part, we assume that platform.dist() will tell the truth.
    # This is the fallback to handle unknowns or exceptions
    OSDIST_DICT = { '/etc/redhat-release': 'RedHat',
                    '/etc/vmware-release': 'VMwareESX' }
    SELINUX_MODE_DICT = { 1: 'enforcing', 0: 'permissive', -1: 'disabled' }

    def __init__(self):
        self.facts = {}
        self.get_platform_facts()
        self.get_distribution_facts()
        self.get_public_ssh_host_keys()
        self.get_selinux_facts()

    def populate(self):
        return self.facts

    # Platform
    # patform.system() can be Linux, Darwin, Java, or Windows
    def get_platform_facts(self):
        self.facts['system'] = platform.system()
        self.facts['kernel'] = platform.release()
        self.facts['machine'] = platform.machine()
        self.facts['python_version'] = platform.python_version()
        self.facts['fqdn'] = socket.getfqdn()
        self.facts['hostname'] = self.facts['fqdn'].split('.')[0]
        if self.facts['machine'] == 'x86_64':
            self.facts['architecture'] = self.facts['machine']
        elif Facts._I386RE.search(self.facts['machine']):
            self.facts['architecture'] = 'i386'
        else:
            self.facts['archtecture'] = self.facts['machine']
        if self.facts['system'] == 'Linux':
            self.get_distribution_facts()

    # platform.dist() is deprecated in 2.6
    # in 2.6 and newer, you should use platform.linux_distribution()
    def get_distribution_facts(self):
        dist = platform.dist()
        self.facts['distribution'] = dist[0].capitalize() or 'NA'
        self.facts['distribution_version'] = dist[1] or 'NA'
        self.facts['distribution_release'] = dist[2] or 'NA'
        # Try to handle the exceptions now ...
        for (path, name) in Facts.OSDIST_DICT.items():
            if os.path.exists(path):
                if self.facts['distribution'] == 'Fedora':
                    pass
                elif name == 'RedHat':
                    data = get_file_content(path)
                    if 'Red Hat' in data:
                        self.facts['distribution'] = name
                    else:
                        self.facts['distribution'] = data.split()[0]
                else:
                    self.facts['distribution'] = name

    def get_public_ssh_host_keys(self):
        dsa = get_file_content('/etc/ssh/ssh_host_dsa_key.pub')
        rsa = get_file_content('/etc/ssh/ssh_host_rsa_key.pub')
        if dsa is None:
            dsa = 'NA'
        else:
            self.facts['ssh_host_key_dsa_public'] = dsa.split()[1]
        if rsa is None:
            rsa = 'NA'
        else:
            self.facts['ssh_host_key_rsa_public'] = rsa.split()[1]

    def get_selinux_facts(self):
        if not HAVE_SELINUX:
            self.facts['selinux'] = False
            return
        self.facts['selinux'] = {}
        if not selinux.is_selinux_enabled():
            self.facts['selinux']['status'] = 'disabled'
        else:
            self.facts['selinux']['status'] = 'enabled'
            self.facts['selinux']['policyvers'] = selinux.security_policyvers()
            (rc, configmode) = selinux.selinux_getenforcemode()
            if rc == 0 and Facts.SELINUX_MODE_DICT.has_key(configmode):
                self.facts['selinux']['config_mode'] = Facts.SELINUX_MODE_DICT[configmode]
            mode = selinux.security_getenforce()
            if Facts.SELINUX_MODE_DICT.has_key(mode):
                self.facts['selinux']['mode'] = Facts.SELINUX_MODE_DICT[mode]
            (rc, policytype) = selinux.selinux_getpolicytype()
            if rc == 0:
                self.facts['selinux']['type'] = policytype

class Hardware(Facts):
    """
    This is a generic Hardware subclass of Facts.  This should be further
    subclassed to implement per platform.  If you subclass this, it
    should define:
    - memfree_mb
    - memtotal_mb
    - swapfree_mb
    - swaptotal_mb
    - processor (a list)
    - processor_cores
    - processor_count

    All subclasses MUST define platform.
    """
    platform = 'Generic'

    def __new__(cls, *arguments, **keyword):
        subclass = cls
        for sc in Hardware.__subclasses__():
            if sc.platform == platform.system():
                subclass = sc
        return super(cls, subclass).__new__(subclass, *arguments, **keyword)

    def __init__(self):
        Facts.__init__(self)

    def populate(self):
        return self.facts

class LinuxHardware(Hardware):
    """
    Linux-specific subclass of Hardware.  Defines memory and CPU facts:
    - memfree_mb
    - memtotal_mb
    - swapfree_mb
    - swaptotal_mb
    - processor (a list)
    - processor_cores
    - processor_count

    In addition, it also defines number of DMI facts.
    """
    platform = 'Linux'
    MEMORY_FACTS = ['MemTotal', 'SwapTotal', 'MemFree', 'SwapFree']
    # DMI bits
    DMI_DICT = { 'form_factor':  '/sys/devices/virtual/dmi/id/chassis_type',
                 'product_name': '/sys/devices/virtual/dmi/id/product_name',
                 'product_serial': '/sys/devices/virtual/dmi/id/product_serial',
                 'product_uuid': '/sys/devices/virtual/dmi/id/product_uuid',
                 'product_version': '/sys/devices/virtual/dmi/id/product_version',
                 'system_vendor': '/sys/devices/virtual/dmi/id/sys_vendor',
                 'bios_date': '/sys/devices/virtual/dmi/id/bios_date',
                 'bios_version': '/sys/devices/virtual/dmi/id/bios_version' }
    # From smolt and DMI spec
    FORM_FACTOR = [ "Unknown", "Other", "Unknown", "Desktop",
                    "Low Profile Desktop", "Pizza Box", "Mini Tower", "Tower",
                    "Portable", "Laptop", "Notebook", "Hand Held", "Docking Station",
                    "All In One", "Sub Notebook", "Space-saving", "Lunch Box",
                    "Main Server Chassis", "Expansion Chassis", "Sub Chassis",
                    "Bus Expansion Chassis", "Peripheral Chassis", "RAID Chassis",
                    "Rack Mount Chassis", "Sealed-case PC", "Multi-system",
                    "CompactPCI", "AdvancedTCA" ]

    def __init__(self):
        Hardware.__init__(self)

    def populate(self):
        self.get_cpu_facts()
        self.get_memory_facts()
        self.get_dmi_facts()
        return self.facts

    def get_memory_facts(self):
        if not os.access("/proc/meminfo", os.R_OK):
            return
        for line in open("/proc/meminfo").readlines():
            data = line.split(":", 1)
            key = data[0]
            if key in LinuxHardware.MEMORY_FACTS:
                val = data[1].strip().split(' ')[0]
                self.facts["%s_mb" % key.lower()] = long(val) / 1024

    def get_cpu_facts(self):
        i = 0
        physid = 0
        sockets = {}
        if not os.access("/proc/cpuinfo", os.R_OK):
            return
        self.facts['processor'] = []
        for line in open("/proc/cpuinfo").readlines():
            data = line.split(":", 1)
            key = data[0].strip()
            if key == 'model name':
                if 'processor' not in self.facts:
                    self.facts['processor'] = []
                self.facts['processor'].append(data[1].strip())
                i += 1
            elif key == 'physical id':
                physid = data[1].strip()
                if physid not in sockets:
                    sockets[physid] = 1
            elif key == 'cpu cores':
                sockets[physid] = int(data[1].strip())
        if len(sockets) > 0:
            self.facts['processor_count'] = len(sockets)
            self.facts['processor_cores'] = reduce(lambda x, y: x + y, sockets.values())
        else:
            self.facts['processor_count'] = i
            self.facts['processor_cores'] = 'NA'

    def get_dmi_facts(self):
        for (key,path) in LinuxHardware.DMI_DICT.items():
            data = get_file_content(path)
            if data is not None:
                if key == 'form_factor':
                    self.facts['form_factor'] = LinuxHardware.FORM_FACTOR[int(data)]
                else:
                    self.facts[key] = data
            else:
                self.facts[key] = 'NA'

class SunOSHardware(Hardware):
    """
    In addition to the generic memory and cpu facts, this also sets
    swap_reserved_mb and swap_allocated_mb that is available from *swap -s*.
    """
    platform = 'SunOS'

    def __init__(self):
        Hardware.__init__(self)

    def populate(self):
        self.get_cpu_facts()
        self.get_memory_facts()
        return self.facts

    def get_cpu_facts(self):
        cmd = subprocess.Popen("/usr/sbin/psrinfo -v", shell=True,
                               stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = cmd.communicate()
        self.facts['processor'] = []
        for line in out.split('\n'):
            if 'processor operates' in line:
                if 'processor' not in self.facts:
                    self.facts['processor'] = []
                self.facts['processor'].append(line.strip())
        self.facts['processor_cores'] = 'NA'
        self.facts['processor_count'] = len(self.facts['processor'])

    def get_memory_facts(self):
        cmd = subprocess.Popen("/usr/sbin/prtconf", shell=False,
                               stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = cmd.communicate()
        for line in out.split('\n'):
            if 'Memory size' in line:
                self.facts['memtotal_mb'] = line.split()[2]
        cmd = subprocess.Popen("/usr/sbin/swap -s", shell=True,
                               stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = cmd.communicate()
        allocated = long(out.split()[1][:-1])
        reserved = long(out.split()[5][:-1])
        used = long(out.split()[8][:-1])
        free = long(out.split()[10][:-1])
        self.facts['swapfree_mb'] = free / 1024
        self.facts['swaptotal_mb'] = (free + used) / 1024
        self.facts['swap_allocated_mb'] = allocated / 1024
        self.facts['swap_reserved_mb'] = reserved / 1024

class FreeBSDHardware(Hardware):
    """
    FreeBSD-specific subclass of Hardware.  Defines memory and CPU facts:
    - memfree_mb
    - memtotal_mb
    - swapfree_mb
    - swaptotal_mb
    - processor (a list)
    - processor_cores
    - processor_count
    """
    platform = 'FreeBSD'
    DMESG_BOOT = '/var/run/dmesg.boot'

    def __init__(self):
        Hardware.__init__(self)

    def populate(self):
        self.get_cpu_facts()
        self.get_memory_facts()
        return self.facts

    def get_cpu_facts(self):
        self.facts['processor'] = []
        cmd = subprocess.Popen("/sbin/sysctl -n hw.ncpu", shell=True,
                               stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = cmd.communicate()
        self.facts['processor_count'] = out.strip()
        for line in open(FreeBSDHardware.DMESG_BOOT).readlines():
            if 'CPU:' in line:
                cpu = re.sub(r'CPU:\s+', r"", line)
                self.facts['processor'].append(cpu.strip())
            if 'Logical CPUs per core' in line:
                self.facts['processor_cores'] = line.split()[4]

    def get_memory_facts(self):
        cmd = subprocess.Popen("/sbin/sysctl vm.stats", shell=True,
                               stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = cmd.communicate()
        for line in out.split('\n'):
            data = line.split()
            if 'vm.stats.vm.v_page_size' in line:
                pagesize = long(data[1])
            if 'vm.stats.vm.v_page_count' in line:
                pagecount = long(data[1])
            if 'vm.stats.vm.v_free_count' in line:
                freecount = long(data[1])
        self.facts['memtotal_mb'] = pagesize * pagecount / 1024 / 1024
        self.facts['memfree_mb'] = pagesize * freecount / 1024 / 1024
        # Get swapinfo.  swapinfo output looks like:
        # Device          1M-blocks     Used    Avail Capacity
        # /dev/ada0p3        314368        0   314368     0%
        # 
        cmd = subprocess.Popen("/usr/sbin/swapinfo -m", shell=True,
                               stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = cmd.communicate()
        lines = out.split('\n')
        if len(lines[-1]) == 0:
            lines.pop()
        data = lines[-1].split()
        self.facts['swaptotal_mb'] = data[1]
        self.facts['swapfree_mb'] = data[3]

class Network(Facts):
    """
    This is a generic Network subclass of Facts.  This should be further
    subclassed to implement per platform.  If you subclass this,
    you must define:
    - interfaces (a list of interface names)
    - interface_<name> dictionary of ipv4, ipv6, and mac address information.

    All subclasses MUST define platform.
    """
    platform = 'Generic'

    def __new__(cls, *arguments, **keyword):
        subclass = cls
        for sc in Network.__subclasses__():
            if sc.platform == platform.system():
                subclass = sc
        return super(cls, subclass).__new__(subclass, *arguments, **keyword)

    def __init__(self):
        Facts.__init__(self)

    def populate(self):
        return self.facts

class LinuxNetwork(Network):
    """
    This is a Linux-specific subclass of Network.  It defines
    - interfaces (a list of interface names)
    - interface_<name> dictionary of ipv4, ipv6, and mac address information.
    """
    platform = 'Linux'
    SIOCGIFCONF   = 0x8912
    SIOCGIFHWADDR = 0x8927

    def __init__(self):
        Network.__init__(self)

    def populate(self):
        self.get_network_facts()
        return self.facts

    # get list of interfaces that are up
    def get_interfaces(self):
        length = 4096
        offset = 32
        step = 32
        if platform.architecture()[0] == '64bit':
            offset = 16
            step = 40
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        names = array.array('B', '\0' * length)
        bytelen = struct.unpack('iL', fcntl.ioctl(
            s.fileno(), LinuxNetwork.SIOCGIFCONF, struct.pack(
                'iL', length, names.buffer_info()[0])
            ))[0]
        return [names.tostring()[i:i+offset].split('\0', 1)[0]
                for i in range(0, bytelen, step)]

    def get_iface_hwaddr(self, iface):
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        info = fcntl.ioctl(s.fileno(), LinuxNetwork.SIOCGIFHWADDR,
                           struct.pack('256s', iface[:15]))
        return ''.join(['%02x:' % ord(char) for char in info[18:24]])[:-1]

    def get_network_facts(self):
        self.facts['interfaces'] = self.get_interfaces()
        for iface in self.facts['interfaces']:
            self.facts[iface] = { 'macaddress': self.get_iface_hwaddr(iface) }
            # This is lame, but there doesn't appear to be a good way
            # to get all addresses for both IPv4 and IPv6.
            cmd = subprocess.Popen("/sbin/ifconfig %s" % iface, shell=True,
                                   stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            out, err = cmd.communicate()
            for line in out.split('\n'):
                is_ipv4 = False
                data = line.split()
                if 'inet addr' in line:
                    if 'ipv4' not in self.facts[iface]:
                        self.facts[iface]['ipv4'] = {}
                    is_ipv4 = True
                    self.facts[iface]['ipv4'] = { 'address': data[1].split(':')[1],
                                                  'netmask': data[-1].split(':')[1] }
                # Slightly different output in net-tools-1.60-134.20120127git
                # Looks like
                # inet 192.168.1.2  netmask 255.255.255.0  broadcast 192.168.1.255
                elif 'inet ' in line:
                    is_ipv4 = True
                    if 'ipv4' not in self.facts[iface]:
                        self.facts[iface]['ipv4'] = {}
                    self.facts[iface]['ipv4'] = { 'address': data[1],
                                                  'netmask': data[3] }
                if is_ipv4:
                    ip = struct.unpack("!L", socket.inet_aton(self.facts[iface]['ipv4']['address']))[0]
                    mask = struct.unpack("!L", socket.inet_aton(self.facts[iface]['ipv4']['netmask']))[0]
                    self.facts[iface]['ipv4']['network'] = socket.inet_ntoa(struct.pack("!L", ip & mask))
                if 'inet6 addr' in line:
                    (ip, prefix) = data[2].split('/')
                    scope = data[3].split(':')[1].lower()
                    if 'ipv6' not in self.facts[iface]:
                        self.facts[iface]['ipv6'] = []
                    self.facts[iface]['ipv6'].append( { 'address': ip,
                                                        'prefix': prefix,
                                                        'scope': scope } )
                # Slightly different output in net-tools-1.60-134.20120127git
                # Looks like
                # inet6 fe80::21b:21ff:fe6e:5baa  prefixlen 64  scopeid 0x20<link>
                elif 'inet6 ' in line:
                    if 'ipv6' not in self.facts[iface]:
                        self.facts[iface]['ipv6'] = []
                    scope = data[5].split('<')[1][:-1]
                    self.facts[iface]['ipv6'].append( { 'address': data[1],
                                                        'prefix': data[3],
                                                        'scope': scope } )

class Virtual(Facts):
    """
    This is a generic Virtual subclass of Facts.  This should be further
    subclassed to implement per platform.  If you subclass this,
    you should define:
    - virtualization_type
    - virtualization_role

    All subclasses MUST define platform.
    """

    def __new__(cls, *arguments, **keyword):
        subclass = cls
        for sc in Virtual.__subclasses__():
            if sc.platform == platform.system():
                subclass = sc
        return super(cls, subclass).__new__(subclass, *arguments, **keyword)

    def __init__(self):
        Facts.__init__(self)

    def populate(self):
        return self.facts

class LinuxVirtual(Virtual):
    """
    This is a Linux-specific subclass of Virtual.  It defines
    - virtualization_type
    - virtualization_role
    """
    platform = 'Linux'

    def __init__(self):
        Virtual.__init__(self)

    def populate(self):
        self.get_virtual_facts()
        return self.facts

    def get_virtual_facts(self):
        if os.path.exists("/proc/xen"):
            self.facts['virtualization_type'] = 'xen'
            self.facts['virtualization_role'] = 'guest'
            if os.path.exists("/proc/xen/capabilities"):
                self.facts['virtualization_role'] = 'host'
        if os.path.exists("/proc/modules"):
            modules = []
            for line in open("/proc/modules").readlines():
                data = line.split(" ", 1)
                modules.append(data[0])
            if 'kvm' in modules:
                self.facts['virtualization_type'] = 'kvm'
                self.facts['virtualization_role'] = 'host'
            elif 'vboxdrv' in modules:
                self.facts['virtualization_type'] = 'virtualbox'
                self.facts['virtualization_role'] = 'host'
            elif 'vboxguest' in modules:
                self.facts['virtualization_type'] = 'virtualbox'
                self.facts['virtualization_role'] = 'guest'
        if 'processor' in self.facts and 'QEMU' in self.facts['processor'][0]:
            self.facts['virtualization_type'] = 'kvm'
            self.facts['virtualization_role'] = 'guest'
        if 'distribution' in self.facts and self.facts['distribution'] == 'VMwareESX':
            self.facts['virtualization_type'] = 'VMware'
            self.facts['virtualization_role'] = 'host'
        # You can spawn a dmidecode process and parse that or infer from devices
        for dev_model in glob.glob('/sys/block/?da/device/vendor'):
            info = open(dev_model).read()
            if 'VMware' in info:
                self.facts['virtualization_type'] = 'VMware'
                self.facts['virtualization_role'] = 'guest'
            elif 'Virtual HD' in info or 'Virtual CD' in info:
                self.facts['virtualization_type'] = 'VirtualPC'
                self.facts['virtualization_role'] = 'guest'

def get_file_content(path):
    data = None
    if os.path.exists(path) and os.access(path, os.R_OK):
        data = open(path).read().strip()
        if len(data) == 0:
            data = None
    return data

def ansible_facts():
    facts = {}
    facts.update(Facts().populate())
    facts.update(Hardware().populate())
    facts.update(Network().populate())
    facts.update(Virtual().populate())
    return facts

# ===========================================

# load config & template variables

if len(sys.argv) == 1:
    sys.exit(1)


argfile = sys.argv[1]
if not os.path.exists(argfile):
    sys.exit(1)

setup_options = open(argfile).read().strip()
try:
   setup_options = json.loads(setup_options)
except:
   list_options = shlex.split(setup_options)
   setup_options = {}
   for opt in list_options:
       (k,v) = opt.split("=")
       setup_options[k]=v   

syslog.openlog('ansible-%s' % os.path.basename(__file__))
syslog.syslog(syslog.LOG_NOTICE, 'Invoked with %s' % setup_options)

# Get some basic facts in case facter or ohai are not installed
for (k, v) in ansible_facts().items():
    setup_options["ansible_%s" % k] = v

# if facter is installed, and we can use --json because
# ruby-json is ALSO installed, include facter data in the JSON

if os.path.exists("/usr/bin/facter"):
   cmd = subprocess.Popen("/usr/bin/facter --json", shell=True,
       stdout=subprocess.PIPE, stderr=subprocess.PIPE)
   out, err = cmd.communicate()
   facter = True
   try:
       facter_ds = json.loads(out)
   except:
       facter = False
   if facter:
       for (k,v) in facter_ds.items():
           setup_options["facter_%s" % k] = v

# ditto for ohai, but just top level string keys
# because it contains a lot of nested stuff we can't use for
# templating w/o making a nicer key for it (TODO)

if os.path.exists("/usr/bin/ohai"):
   cmd = subprocess.Popen("/usr/bin/ohai", shell=True,
       stdout=subprocess.PIPE, stderr=subprocess.PIPE)
   out, err = cmd.communicate()
   ohai = True
   try:
       ohai_ds = json.loads(out)
   except:
       ohai = False
   if ohai:
       for (k,v) in ohai_ds.items():
           if type(v) == str or type(v) == unicode:
               k2 = "ohai_%s" % k
               setup_options[k2] = v

setup_result = {}
setup_result['ansible_facts'] = setup_options

# hack to keep --verbose from showing all the setup module results
setup_result['verbose_override'] = True

print json.dumps(setup_result)