#!/usr/bin/env 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


client = None
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 = 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 == None:
        client.connect(opt_host, username=opt_user, password=opt_pass, timeout=opt_timeout)
    else:
        client.connect(opt_host, username=opt_user, pkey=opt_key, 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]
    if filename in sftp.listdir("%s/%s" % (working_dir, opt_put_remote)):
        return True
    else:
        return False


def create_testfile(paths):
    path = paths["local_put_path"]
    if not os.path.isfile(path):
        os.system('echo "This is a test by Check_MK" > %s' % path)


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 != 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 != 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 != 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()
