#!/usr/bin/python
# -*- encoding: utf-8; py-indent-offset: 4 -*-
# +------------------------------------------------------------------+
# |             ____ _               _        __  __ _  __           |
# |            / ___| |__   ___  ___| | __   |  \/  | |/ /           |
# |           | |   | '_ \ / _ \/ __| |/ /   | |\/| | ' /            |
# |           | |___| | | |  __/ (__|   <    | |  | | . \            |
# |            \____|_| |_|\___|\___|_|\_\___|_|  |_|_|\_\           |
# |                                                                  |
# | Copyright Mathias Kettner 2014             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 os
import sys
import time
import getopt
import paramiko

sftp = None


def usage():
    sys.stderr.write("""
USAGE: check_sftp [OPTIONS] HOST

OPTIONS:
  --host HOST            SFTP server address
  --user USER            Username for sftp login
  --secret SECRET        Secret/Password for sftp login
  --private-key KEY      Private Key for sftp login
  --port PORT            Alternative port number (default is 22 for the connection)
  --get-remote FILE      Path to the file which to pull from SFTP server (e.g.
                         /tmp/testfile.txt)
  --get-local PATH       Path to store the pulled file locally (e.g. $OMD_ROOT/tmp/)
  --put-local FILE       Path to the file to push to the sftp server. See above for example
  --put-remote PATH      Path to save the pushed file (e.g. /tmp/)
  --get-timestamp PATH   Path to the file for getting the timestamp of this file
  --timeout SECONDS      Set timeout for connection (default is 10 seconds)
  --verbose              Output some more detailed information
  -h, --help             Show this help message and exit
    """)
    sys.exit(1)


opt_host = None
opt_user = None
opt_pass = None
opt_key = None
opt_port = 22
opt_get_remote = None
opt_get_local = None
opt_put_local = None
opt_put_remote = None
opt_timestamp = None
opt_timeout = 10.0
opt_verbose = False

short_options = 'hv'
long_options = [
    'host=',
    'user=',
    'secret=',
    'privat-key',
    'port=',
    'get-remote=',
    'get-local=',
    'put-local=',
    'put-remote=',
    'get-timestamp=',
    'verbose',
    'help',
    'timeout=',
]

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 opt, arg in opts:
    if opt in ['-h', 'help']:
        usage()
    elif opt in ['--host']:
        opt_host = arg
    elif opt in ['--user']:
        opt_user = arg
    elif opt in ['--secret']:
        opt_pass = arg
    elif opt in ['--private-key']:
        opt_pass = arg
    elif opt in ['--port']:
        opt_port = int(arg)
    elif opt in ['--timeout']:
        opt_timeout = float(arg)
    elif opt in ['--put-local']:
        opt_put_local = arg
    elif opt in ['--put-remote']:
        opt_put_remote = arg
    elif opt in ['--get-local']:
        opt_get_local = arg
    elif opt in ['--get-remote']:
        opt_get_remote = arg
    elif opt in ['--get-timestamp']:
        opt_timestamp = arg
    elif opt in ['-v', '--verbose']:
        opt_verbose = True


def connection():
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    if opt_key is None:
        client.connect(opt_host,
                       username=opt_user,
                       password=opt_pass,
                       port=opt_port,
                       timeout=opt_timeout)
    else:
        client.connect(opt_host,
                       username=opt_user,
                       pkey=opt_key,
                       port=opt_port,
                       timeout=opt_timeout)
    return client


def get_paths(omd_root, working_dir):
    paths = {}
    if opt_put_local:
        put_filename = opt_put_local.split("/")[-1]
        paths["put_filename"] = put_filename
        paths["local_put_path"] = "%s/%s" % (omd_root, opt_put_local)
        if len(opt_put_remote) > 0:
            paths["remote_put_path"] = "%s/%s/%s" % (working_dir, opt_put_remote, put_filename)
        else:
            paths["remote_put_path"] = "%s/%s" % (working_dir, put_filename)

    if opt_get_remote:
        get_filename = opt_get_remote.split("/")[-1]
        paths["get_filename"] = get_filename
        paths["remote_get_path"] = "%s/%s" % (working_dir, opt_get_remote)
        if len(opt_get_local) > 0:
            paths["local_get_path"] = "%s/%s/%s" % (omd_root, opt_get_local, get_filename)
        else:
            paths["local_get_path"] = "%s/%s" % (omd_root, get_filename)

    if opt_timestamp:
        paths["timestamp_filename"] = opt_timestamp.split("/")[-1]
        paths["timestamp_path"] = "%s/%s" % (working_dir, opt_timestamp)

    return paths


def file_available(working_dir):
    filename = opt_put_local.split("/")[-1]
    return filename in sftp.listdir("%s/%s" % (working_dir, opt_put_remote))


def create_testfile(paths):
    path = paths["local_put_path"]
    if not os.path.isfile(path):
        with open(path, "w") as f:
            f.write("This is a test by Check_MK\n")


def put_file(paths):
    sftp.put(paths["local_put_path"], \
             paths["remote_put_path"])


def get_file(paths):
    sftp.get(paths["remote_get_path"], \
             paths["local_get_path"])


def get_timestamp(paths):
    return sftp.stat(paths["timestamp_path"])


def write_out(state, message):
    state_readable = ['OK', 'WARN', 'CRIT', 'UNKNOWN'][state]
    sys.stdout.write('%s - %s\n' % (state_readable, message))
    sys.exit(state)


def main():
    global sftp
    messages = []
    states = []
    try:  # Establish connection
        client = connection()
        sftp = client.open_sftp()
        messages.append("Login successful")
        states.append(0)
    except:
        if opt_verbose:
            raise
        write_out(2, "Connection failed!")

    # Let's prepare for some other tests...
    omd_root = os.getenv("OMD_ROOT")
    sftp.chdir(".")
    working_dir = sftp.getcwd()
    paths = get_paths(omd_root, working_dir)
    testfile_remote = True

    # .. and eventually execute them!
    try:  # Put a file to the server
        if opt_put_local is not None:
            create_testfile(paths)
            testfile_remote = file_available(working_dir)
            put_file(paths)
            states.append(0)
            messages.append("Successfully put file to SFTP server")
    except:
        if opt_verbose:
            raise
        states.append(2)
        messages.append("Could not put file to SFTP server! (!!)")

    try:  # Get a file from the server
        if opt_get_remote is not None:
            get_file(paths)
            states.append(0)
            messages.append("Successfully got file from SFTP server")
    except:
        if opt_verbose:
            raise
        states.append(2)
        messages.append("Could not get file from SFTP server! (!!)")

    try:  # Get timestamp of a remote file
        if opt_timestamp is not None:
            file_stats = get_timestamp(paths)
            states.append(0)
            messages.append("Timestamp of %s is: %s" % \
                            (paths["timestamp_filename"], \
                             time.ctime(file_stats.st_mtime)))
    except:
        if opt_verbose:
            raise
        states.append(2)
        messages.append("Could not get timestamp of file! (!!)")

    # Remove useless files
    if not testfile_remote:
        sftp.remove(paths["remote_put_path"])

    write_out(max(states), ", ".join(messages))


main()
