#!/usr/bin/python
# -*- encoding: utf-8; py-indent-offset: 4 -*-
# +------------------------------------------------------------------+
# |             ____ _               _        __  __ _  __           |
# |            / ___| |__   ___  ___| | __   |  \/  | |/ /           |
# |           | |   | '_ \ / _ \/ __| |/ /   | |\/| | ' /            |
# |           | |___| | | |  __/ (__|   <    | |  | | . \            |
# |            \____|_| |_|\___|\___|_|\_\___|_|  |_|_|\_\           |
# |                                                                  |
# | Copyright Mathias Kettner 2015             mk@mathias-kettner.de |
# +------------------------------------------------------------------+
#
# This file is part of Check_MK.
# The official homepage is at http://mathias-kettner.de/check_mk.
#
# check_mk 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 in version 2.  check_mk is  distributed
# in the hope that it will be useful, but WITHOUT ANY WARRANTY;  with-
# out even the implied warranty of  MERCHANTABILITY  or  FITNESS FOR A
# PARTICULAR PURPOSE. See the  GNU General Public License for more de-
# tails. You should have  received  a copy of the  GNU  General Public
# License along with GNU Make; see the file  COPYING.  If  not,  write
# to the Free Software Foundation, Inc., 51 Franklin St,  Fifth Floor,
# Boston, MA 02110-1301 USA.

import getopt, sys, socket, traceback, re, pprint
import snap7
from snap7.util import *
from snap7.snap7types import S7AreaPE, S7AreaPA, S7AreaMK, S7AreaDB, \
                             S7AreaCT, S7AreaTM

def usage():
    sys.stderr.write("""Check_MK Siemens PLC Agent

USAGE: agent_siemens_plc [OPTIONS] HOSTSPECS...
       agent_siemens_plc -h

HOSTSPECS:
  A HOSTSPEC specifies the hosts to contact and the data to fetch
  from each host. A hostspec is built of minimum 5 ";"
  separated items, which are:

  HOST_NAME                     Logical name of the PLC
  HOST_ADDRESS                  Host name or IP address of the PLC
  RACK
  SLOT
  PORT                          The TCP port to communicate with
  VALUES                        One or several VALUES as defined below.
                                The values themselfs are separated by ";"

VALUES:
  A value is specified by the following single data fields, which are
  concatenated by a ",":

    AREA[:DB_NUMBER]            Identifier of the memory area to fetch (db, input,
                                output, merker, timer or counter), plus the optional
                                numeric identifier of the DB separeated by a ":".
    ADDRESS                     Memory address to read
    DATATYPE                    The datatype of the value to read
    VALUETYPE                   The logical type of the value
    IDENT                       An identifier of your choice. This identifier
                                is used by the Check_MK checks to access
                                and identify the single values. The identifier
                                needs to be unique within a group of VALUETYPES.

OPTIONS:
  -h, --help                    Show this help message and exit
  -t, --timeout SEC             Set the network timeout to <SEC> seconds.
                                Default is 10 seconds. Note: the timeout is not
                                applied to the whole check, instead it is used for
                                each network connect.
  --debug                       Debug mode: let Python exceptions come through
""")

short_options = 'h:t:d'
long_options  = [
    'help', 'timeout=', 'debug'
]

devices           = []
opt_debug         = False
opt_timeout       = 10

try:
    opts, args = getopt.getopt(sys.argv[1:], short_options, long_options)
except getopt.GetoptError, err:
    sys.stderr.write("%s\n" % err)
    sys.exit(1)

for o,a in opts:
    if o in [ '--debug' ]:
        opt_debug = True
    elif o in [ '-t', '--timeout' ]:
        opt_timeout = int(a)
    elif o in [ '-h', '--help' ]:
        usage()
        sys.exit(0)

if not args:
    sys.stderr.write("ERROR: You missed to provide the needed arguments.\n")
    usage()
    sys.exit(1)

for arg in args:
    parts = arg.split(';')
    if len(parts) < 6:
        sys.stderr.write("ERROR: Not enough arguments: %s\n" % arg)
        usage()
        sys.exit(1)

    values = []
    for spec in parts[5:]:
        p = spec.split(',')
        if len(p) != 5:
            sys.stderr.write("ERROR: Invalid value specified: %s\n" % spec)
            usage()
            sys.exit(1)

        if ':' in p[0]:
            area_name, db_number = p[0].split(':')
            area = (area_name, int(db_number))
        elif p[0] in ["merker", "input", "output", "counter", "timer"]:
            area = (p[0], None)
        else:
            area = ("db", int(p[0]))

        byte, bit = map(int, p[1].split('.'))

        if ':' in p[2]:
            typename, size = p[2].split(':')
            datatype = typename, int(size)
        else:
            datatype = p[2]

        # area, address, datatype, valuetype, ident
        values.append((area, (byte, bit), datatype, p[3], p[4]))

    devices.append({
        'host_name'     : parts[0],
        'host_address'  : parts[1],
        'rack'          : int(parts[2]),
        'slot'          : int(parts[3]),
        'port'          : int(parts[4]),
        'values'        : values,
    })

socket.setdefaulttimeout(opt_timeout)


def get_dint(_bytearray, byte_index):
    """
    Get int value from bytearray.

    double int are represented in four bytes
    """
    byte3 = _bytearray[byte_index + 3]
    byte2 = _bytearray[byte_index + 2]
    byte1 = _bytearray[byte_index + 1]
    byte0 = _bytearray[byte_index]
    return byte3 + (byte2 << 8) + (byte1 << 16) + (byte0 << 32)


def area_name_to_area_id(name):
    return {
        'db'      : S7AreaDB,
        'input'   : S7AreaPE,
        'output'  : S7AreaPA,
        'merker'  : S7AreaMK,
        'timer'   : S7AreaTM,
        'counter' : S7AreaCT,
    }[name]


datatypes = {
    # type-name   size(bytes) parse-function
    # A size of None means the size is provided by configuration
    'dint':      (4,          lambda data, offset, size, bit: get_dint(data, offset)),
    'real':      (8,          lambda data, offset, size, bit: get_real(data, offset)),
    'bit':       (1,          lambda data, offset, size, bit: get_bool(data, offset, bit)),
    # str currently handles "zeichen" (character?) formated strings. For byte coded strings
    # we would have to use get_string(data, offset-1)) from snap7.utils
    'str':       (None,       lambda data, offset, size, bit: data[offset:offset+size]),
}

unhandled_error = False
for device in devices:
    try:
        client = snap7.client.Client()
        client.connect(device['host_address'], device['rack'],
                       device['slot'], device['port'])

        sys.stdout.write("<<<siemens_plc_cpu_state>>>\n")
        sys.stdout.write(client.get_cpu_state()+"\n")

        sys.stdout.write("<<<siemens_plc>>>\n")
        # We want to have a minimum number of reads. We try to only use
        # a single read and detect the memory area to fetch dynamically
        # based on the configured values
        addresses = {}
        start_address = None
        end_address   = None
        for area, (byte, bit), datatype, valuetype, ident in device['values']:
            if type(datatype) == tuple:
                size = datatype[1]
            else:
                size = datatypes[datatype][0]
            addresses.setdefault(area, [None, None])
            start_address, end_address = addresses[area]

            if start_address == None or byte < start_address:
                addresses[area][0] = byte

            end = byte + size
            if end_address == None or end > end_address:
                addresses[area][1] = end

        # Now fetch the data from each db number
        data = {}
        for (area_name, db_number), (start, end) in addresses.items():
            area_id = area_name_to_area_id(area_name)
            data[(area_name, db_number)] = client.read_area(area_id, db_number, start, size=end-start)

        # Now loop all values to be fetched and extract the data
        # from the bytes fetched above
        for (area_name, db_number), (byte, bit), datatype, valuetype, ident in device['values']:
            if type(datatype) == tuple:
                typename, size = datatype
                parse_func = datatypes[typename][1]
            else:
                size, parse_func = datatypes[datatype]

            start, end = addresses[(area_name, db_number)]
            fetched_data = data[(area_name, db_number)]

            value = parse_func(fetched_data, byte-start, size, bit)
            sys.stdout.write("%s %s %s %s\n" % (device['host_name'], valuetype, ident, value))
    except:
        sys.stderr.write('%s: Unhandled error: %s' % (device['host_name'], traceback.format_exc()))
        if opt_debug:
            raise
        unhandled_error = True

sys.exit(unhandled_error and 1 or 0)
