#!/usr/bin/env python
# Copyright (C) 2013  Lukas Rist <glaslos@gmail.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.
#
# 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 gevent.monkey
gevent.monkey.patch_all()

import logging
import os
import argparse
import sys
import pwd
import grp
import platform
import ast
import inspect
from ConfigParser import ConfigParser

import gevent
from lxml import etree

import conpot
import conpot.core as conpot_core

from conpot.core.loggers.log_worker import LogWorker
from conpot.protocols.snmp.snmp_server import SNMPServer
from conpot.protocols.modbus.modbus_server import ModbusServer
from conpot.protocols.s7comm.s7_server import S7Server
from conpot.protocols.http.web_server import HTTPServer
from conpot.helpers import fix_sslwrap
from conpot.protocols.kamstrup.meter_protocol.kamstrup_server import KamstrupServer
from conpot.protocols.kamstrup.management_protocol.kamstrup_management_server import KamstrupManagementServer

from conpot.emulators.proxy import Proxy
from conpot.utils import ext_ip


logger = logging.getLogger()
package_directory = os.path.dirname(os.path.abspath(conpot.__file__))


def logo():
    print """
                       _
   ___ ___ ___ ___ ___| |_
  |  _| . |   | . | . |  _|
  |___|___|_|_|  _|___|_|
              |_|

  Version {0}
  Glastopf Project
""".format(conpot.__version__)


def on_unhandled_greenlet_exception(dead_greenlet):
    logger.error('Stopping because {0} died: {1}'.format(dead_greenlet, dead_greenlet.exception))
    sys.exit(1)

def setup_logging(log_file, verbose):
    if verbose:
        log_level = logging.DEBUG
    else:
        log_level = logging.INFO

    log_format = logging.Formatter('%(asctime)-15s %(message)s')
    console_log = logging.StreamHandler()
    console_log.setLevel(log_level)
    console_log.setFormatter(log_format)

    logger.setLevel(log_level)
    file_log = logging.FileHandler(log_file)
    file_log.setFormatter(log_format)
    file_log.setLevel(log_level)

    root_logger = logging.getLogger()
    root_logger.addHandler(console_log)
    root_logger.addHandler(file_log)


def drop_privileges(uid_name='nobody', gid_name='nogroup'):
    wanted_uid = pwd.getpwnam(uid_name)[2]
    # special handling for os x. (getgrname has trouble with gid below 0)
    if platform.mac_ver()[0] and platform.mac_ver()[0] < float('10.9'):
        wanted_gid = -2
    else:
        wanted_gid = grp.getgrnam(gid_name)[2]

    os.setgid(wanted_gid)
    os.setuid(wanted_uid)

    new_uid_name = pwd.getpwuid(os.getuid())[0]
    new_gid_name = grp.getgrgid(os.getgid())[0]

    logger.info("Privileges dropped, running as {0}/{1}.".format(new_uid_name, new_gid_name))


def validate_template(xml_file, xsd_file):
    xml_schema = etree.parse(xsd_file);
    xsd = etree.XMLSchema(xml_schema)
    xml = etree.parse(xml_file)
    xsd.validate(xml)
    if xsd.error_log:
        logger.error('Error parsing XML template: {0}'.format(xsd.error_log))
        sys.exit(1)


def main():
    logo()

    parser = argparse.ArgumentParser()

    parser.add_argument("-t", "--template",
                        help="Name of one of the supplied templates, or the full path to a custom template.",
                        default=''
                        )
    parser.add_argument("-c", "--config",
                        help="The configuration file to use",
                        default=os.path.join(package_directory, 'conpot.cfg'),
                        metavar="config.cfg"
                        )
    parser.add_argument("-l", "--logfile",
                        help="The logfile to use",
                        default="conpot.log"
                        )
    parser.add_argument("-a", "--raw_mib",
                        help="Path to raw MIB files."
                             "(will automatically get compiled by build-pysnmp-mib)",
                        action='append',
                        default=[]
                        )
    parser.add_argument("-m", "--mibpaths",
                        action='append',
                        help="Path to compiled PySNMP MIB files."
                             "(must be compiled with build-pysnmp-mib)",
                        default=[]
                        )
    parser.add_argument('-v', '--verbose', action='store_true', default=False, help='Logs debug messages.')
    args = parser.parse_args()

    setup_logging(args.logfile, args.verbose)

    config = ConfigParser()
    if not os.path.isfile(args.config):
        args.config = os.path.join(package_directory, 'conpot.cfg')
        logger.info('No conpot.cfg found in current directory, using default configuration: {0}'.format(args.config))
    config.read(args.config)

    # No template specified
    if not args.template:

        available_templates = os.listdir(os.path.join(package_directory, 'templates'))

        print "--------------------------------------------------"
        print " Available templates:"
        print "--------------------------------------------------\n"

        for folder in available_templates:
            template_xml = os.path.join(package_directory, 'templates', folder, 'template.xml')
            if os.path.isfile(template_xml):
                template_unit = template_vendor = template_description = template_protocols = template_creator = 'N/A'
                dom_template = etree.parse(template_xml)
                template_details = dom_template.xpath('//core/template/*')
                if template_details:

                    # retrieve all template details
                    for entity in template_details:

                        if entity.attrib['name'] == 'unit':
                            template_unit = entity.text

                        elif entity.attrib['name'] == 'vendor':
                            template_vendor = entity.text

                        elif entity.attrib['name'] == 'description':
                            template_description = entity.text

                        elif entity.attrib['name'] == 'protocols':
                            template_protocols = entity.text

                        elif entity.attrib['name'] == 'creator':
                            template_creator = entity.text

                    print "   --template {0}".format(folder)
                    print "       Unit:        {0} - {1}".format(template_vendor, template_unit)
                    print "       Desc:        {0}".format(template_description)
                    print "       Protocols:   {0}".format(template_protocols)
                    print "       Created by:  {0}\n".format(template_creator)

        sys.exit(0)

    #  Custom template supplied
    if os.path.exists(os.path.join(args.template, 'template.xml')):
        root_template_directory = args.template
    # Check if the template name can be in the default templates directory
    elif os.path.isfile(os.path.join(package_directory, 'templates', args.template, 'template.xml')):
        root_template_directory = os.path.join(package_directory, 'templates', args.template)
    else:
            logger.error('Template not found: {0}'.format(args.template))
            sys.exit(1)

    # Check if the configuration file exists..
    if not os.path.isfile(args.config):
        logger.error('Config not found: {0}'.format(args.config))
        sys.exit(1)

    logger.info('Starting Conpot using template: {0}'.format(root_template_directory))
    logger.info('Starting Conpot using configuration found in: {0}'.format(args.config))

    servers = list()

    template_base = os.path.join(root_template_directory, 'template.xml')
    if os.path.isfile(template_base):
        validate_template(template_base, os.path.join(package_directory, 'tests/template_schemas/core.xsd'))
        dom_base = etree.parse(template_base)
    else:
        logger.error('Could not access template configuration')
        sys.exit(1)

    session_manager = conpot_core.get_sessionManager()
    session_manager.initialize_databus(template_base)

    public_ip = None
    if config.getboolean('fetch_public_ip', 'enabled'):
        public_ip = ext_ip.get_ext_ip(config)

    protocol_instance_mapping = (
        ('modbus', ModbusServer),
        ('s7comm', S7Server),
        ('kamstrup_meter', KamstrupServer),
        ('kamstrup_management', KamstrupManagementServer),
        ('http', HTTPServer),
        ('snmp', SNMPServer)
    )

    for protocol in protocol_instance_mapping:
        protocol_name, server_class = protocol
        protocol_template = os.path.join(root_template_directory, protocol_name, '{0}.xml'.format(protocol_name))
        if os.path.isfile(protocol_template):
            xsd_file = os.path.join(os.path.dirname(inspect.getfile(server_class)), '{0}.xsd'.format(protocol_name))
            validate_template(protocol_template, xsd_file)
            dom_protocol = etree.parse(protocol_template)
            if dom_protocol.xpath('//{0}'.format(protocol_name)):
                if ast.literal_eval(dom_protocol.xpath('//{0}/@enabled'.format(protocol_name))[0]):
                    host = dom_protocol.xpath('//{0}/@host'.format(protocol_name))[0]
                    port = ast.literal_eval(dom_protocol.xpath('//{0}/@port'.format(protocol_name))[0])
                    server = server_class(protocol_template, root_template_directory, args)
                    greenlet = gevent.spawn(server.start, host, port)
                    greenlet.link_exception(on_unhandled_greenlet_exception)
                    servers.append(server)
                    logger.info('Found and enabled {0} protocol.'.format(protocol))
            else:
                logger.info('{0} available but disabled by configuration.'.format(protocol_name))
        else:
            logger.debug('No {0} template found. Service will remain unconfigured/stopped.'.format(protocol_name))

    log_worker = LogWorker(config, dom_base, session_manager, public_ip)
    greenlet = gevent.spawn(log_worker.start)
    greenlet.link_exception(on_unhandled_greenlet_exception)

    # TODO: Make proxy fit into protocol_instance_mapping
    template_proxy = os.path.join(root_template_directory, 'proxy', 'proxy.xml')
    if os.path.isfile(template_proxy):
        xsd_file = os.path.join(os.path.dirname(inspect.getfile(Proxy)), 'proxy.xsd')
        validate_template(template_proxy, xsd_file)
        dom_proxy = etree.parse(template_proxy)
        if dom_proxy.xpath('//proxies'):
            if ast.literal_eval(dom_proxy.xpath('//proxies/@enabled')[0]):
                proxies = dom_proxy.xpath('//proxies/*')
                for p in proxies:
                    name = p.attrib['name']
                    host = p.attrib['host']
                    keyfile = None
                    certfile = None
                    if 'keyfile' in p.attrib and 'certfile' in p.attrib:
                        keyfile = p.attrib['keyfile']
                        certfile = p.attrib['certfile']

                        # if path is absolute we assert that the cert and key is located in
                        # the templates ssl standard location

                        if not os.path.isabs(keyfile):
                            keyfile = os.path.join(os.path.dirname(root_template_directory), 'ssl', keyfile)
                            certfile = os.path.join(os.path.dirname(root_template_directory), 'ssl', certfile)
                    port = ast.literal_eval(p.attrib['port'])
                    proxy_host = p.xpath('./proxy_host/text()')[0]
                    proxy_port = ast.literal_eval(p.xpath('./proxy_port/text()')[0])
                    decoder = p.xpath('./decoder/text()')
                    if len(decoder) > 0:
                        decoder = decoder[0]
                    else:
                        decoder = None
                    proxy_instance = Proxy(name, proxy_host, proxy_port, decoder, keyfile, certfile)
                    proxy_server = proxy_instance.get_server(host, port)
                    servers.append(proxy_instance)
                    proxy_greenlet = gevent.spawn(proxy_server.start)
                    proxy_greenlet.link_exception(on_unhandled_greenlet_exception)
            else:
                logger.info('Proxy available but disabled by template.')
    else:
            logger.info('No proxy template found. Service will remain unconfigured/stopped.')

    # Wait for the services to bind ports before dropping privileges
    gevent.sleep(5)
    # Only drop if running as root
    if os.getuid() == 0:
        drop_privileges()
    try:
        if len(servers) > 0:
            gevent.wait()
    except KeyboardInterrupt:
        logging.info('Stopping Conpot')
        for server in servers:
            server.stop()

if __name__ == "__main__":
    fix_sslwrap()
    main()
