Merge pull request #1117 from DataDog/arthur/ssh-check

Arthur/ssh check

checks.d/ssh_check.py

@@ -0,0 +1,93 @@
+# stdlib
+import time
+import socket
+# 3p
+import paramiko
+from collections import namedtuple
+# project
+from checks import AgentCheck
+
+class CheckSSH(AgentCheck):
+
+    OPTIONS = [
+        ('host', True, None, str),
+        ('port', False, 22, int),
+        ('username', True, None, str),
+        ('password', False, None, str),
+        ('private_key_file', False, None, str),
+        ('sftp_check', False, True, bool),
+        ('add_missing_keys', False, False, bool),
+    ]
+
+    Config = namedtuple('Config', [
+                'host',
+                'port',
+                'username',
+                'password',
+                'private_key_file',
+                'sftp_check',
+                'add_missing_keys',
+            ]
+            )
+    def _load_conf(self, instance):
+        params = []
+        for option, required, default, expected_type in self.OPTIONS:
+            value = instance.get(option)
+            if required and (not value or type(value)) != expected_type :
+                raise Exception("Please specify a valid {0}".format(option))
+
+            if value is None or type(value) != expected_type:
+                self.log.debug("Bad or missing value for {0} parameter. Using default".format(option))
+                value = default
+
+            params.append(value)
+        return self.Config._make(params)
+
+    def check(self, instance):
+        conf = self._load_conf(instance)
+
+        try:
+            private_key = paramiko.RSAKey.from_private_key_file (conf.private_key_file)
+        except Exception:
+            self.warning("Private key could not be found")
+            private_key = None
+
+        client = paramiko.SSHClient()
+        if conf.add_missing_keys:
+            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        client.load_system_host_keys()
+
+        exception_message = None
+        #Service Availability to check status of SSH
+        try:
+            client.connect(conf.host, port=conf.port, username=conf.username, password=conf.password, pkey=private_key)
+            self.service_check('ssh.can_connect', AgentCheck.OK, message=exception_message)
+
+        except Exception as e:
+            exception_message = str(e)
+            status = AgentCheck.CRITICAL
+            self.service_check('ssh.can_connect', status, message=exception_message)
+            if conf.sftp_check:
+                self.service_check('sftp.can_connect', status, message=exception_message)
+            raise Exception (e)
+
+        #Service Availability to check status of SFTP
+        if conf.sftp_check:
+            try:
+                sftp = client.open_sftp()
+                #Check response time of SFTP
+                start_time = time.time()
+                result = sftp.listdir('.')
+                status = AgentCheck.OK
+                end_time = time.time()
+                time_taken = end_time - start_time
+                self.gauge('sftp.response_time', time_taken)
+
+            except Exception as e:
+                exception_message = str(e)
+                status = AgentCheck.CRITICAL
+
+            if exception_message is None:
+                exception_message = "No errors occured"
+
+            self.service_check('sftp.can_connect', status, message=exception_message)

conf.d/ssh_check.yaml.example

@@ -0,0 +1,11 @@
+init_config:
+
+instances:
+
+    - host: localhost             #required, must be filled
+      port: 22                    #optional, leaving blank defaults to port 22
+      username: test              #required, must be filled
+      password: abcd              #optional
+      sftp_check: True            #optional, leaving blank defaults to True
+      private_key_file:           #optional, file path to private key
+      add_missing_keys: True      #optional, leaving blank defaults to False

requirements.txt

@@ -18,3 +18,4 @@ ntplib
 httplib2
 kafka-python==0.9.0-9bed11db98387c0d9e456528130b330631dc50af
 requests
+paramiko

tests/test_ssh_check.py

@@ -0,0 +1,56 @@
+import unittest
+from tests.common import load_check
+from checks import AgentCheck
+
+class SshTestCase(unittest.TestCase):
+
+    def test_ssh(self):
+
+        config = {
+            'instances': [{
+                'host': 'sdf.org',
+                'port': 22,
+                'username': 'datadog01',
+                'password': 'abcd',
+                'sftp_check': False,
+                'private_key_file': '',
+                'add_missing_keys': True
+            },
+            {
+                'host': 'sdf.org',
+                'port': 22,
+                'username': 'wrongusername',
+                'password': 'wrongpassword',
+                'sftp_check': False,
+                'private_key_file': '',
+                'add_missing_keys': True
+            },
+            {
+                'host': 'wronghost',
+                'port': 22,
+                'username': 'datadog01',
+                'password': 'abcd',
+                'sftp_check': False,
+                'private_key_file': '',
+                'add_missing_keys': True
+            },
+            ]
+        }
+
+        agentConfig = {}
+        self.check = load_check('ssh_check', config, agentConfig)
+
+        #Testing that connection will work
+        self.check.check(config['instances'][0])
+
+        service = self.check.get_service_checks()
+        self.assertEqual(service[0].get('status'), AgentCheck.OK)
+        self.assertEqual(service[0].get('message'), None)
+
+        #Testing that bad authentication will raise exception
+        self.assertRaises(Exception, self.check.check, config['instances'][1])
+        #Testing that bad hostname will raise exception
+        self.assertRaises(Exception, self.check.check, config['instances'][2])
+        service_fail = self.check.get_service_checks()
+        #Check failure status
+        self.assertEqual(service_fail[0].get('status'), AgentCheck.CRITICAL)