"""Written for Python 3.5, works on 2.7.3 +

Pip install the following moduels:
psycopg2==2.7.1
PyYAML==3.12

Copyright (C) 2017 William Christensen

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 version 2.

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.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

"""

import psycopg2
import yaml
from collections import defaultdict

## Constants
DATABASE_CONFIG = '/etc/foreman/database.yml' # If using Satellite, change foreman to satellite.
OUTPUTFILE_NAME = "inventory.txt"

def main():
    """Reads the database to pull clients that have checked into puppet within the last day.
    Returns a CSV to be parsed of the results."""

    database = ''
    username = ''
    password = ''

    try:
        dbfile = open(DATABASE_CONFIG)
        dataMap = yaml.safe_load(dbfile)

        database = dataMap['production']['database']
        username = dataMap['production']['username']
        password = dataMap['production']['password']

    except:
        print("Failed reading the database config file for The Foreman. Please edit the script for the proper name and location.")
        return

    connection = None
    cur = None
    hosts = None

    try:
        connection = psycopg2.connect("host='localhost' dbname='"+database+"' user='"+username+"' password='"+password+"'") # Connection to database
        cur = connection.cursor() # Cursor object to the database.
    except:
        print("Connection with database failed.")
        return

    try:
        # Select all hosts that have reported in the last day.
        cur.execute("SELECT name from hosts WHERE last_report > (current_date - 1) ORDER BY name;")
        hosts = cur.fetchall()
    except:
        print("Could not perform select statement")
        return

    try:

        groups = defaultdict(list)

        for host in hosts:
            endLoc = host[0].index(".")
            # In the event systems wish to be classified by hostname, add the logic here.
            ### <Change for your naming scheme here for grouping of systems> ###
            if(endLoc > 0 or endLoc != None):
                print(host[0][0:endLoc])
                groups[host[0][0:3]].append(host[0][0:endLoc] + "\n")
            else:
                print(host[0])
                groups[host[0][0:3]].append(host[0] + "\n")
             ### </Change for your naming scheme here for grouping of systems> ###

        output = open(OUTPUTFILE_NAME,"w")

        for key in groups.keys():
            output.write("["+ key + "]" + "\n") # print the group name by first 3 letters
            for host in groups[key]:
                output.write(host) # print the host names from the list.
            output.write("\n")

        output.close()

    except:
        print("hosts could not be populated or found.")
        return


if __name__ == '__main__':
    main()