TLS settings and tests

Author: technige

neo4j/v1/__init__.py

@@ -18,5 +18,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .constants import *
 from .session import *
 from .typesystem import *

neo4j/v1/compat.py

@@ -90,19 +90,3 @@ def perf_counter():
     from urllib.parse import urlparse
 except ImportError:
     from urlparse import urlparse
-
-
-try:
-    from ssl import SSLContext, PROTOCOL_SSLv23, OP_NO_SSLv2, HAS_SNI
-except ImportError:
-    from ssl import wrap_socket, PROTOCOL_SSLv23
-
-    def secure_socket(s, host):
-        return wrap_socket(s, ssl_version=PROTOCOL_SSLv23)
-
-else:
-
-    def secure_socket(s, host):
-        ssl_context = SSLContext(PROTOCOL_SSLv23)
-        ssl_context.options |= OP_NO_SSLv2
-        return ssl_context.wrap_socket(s, server_hostname=host if HAS_SNI else None)

neo4j/v1/connection.py

@@ -21,25 +21,24 @@
 
 from __future__ import division
 
+from base64 import b64encode
 from collections import deque
 from io import BytesIO
 import logging
-from os import environ
+from os import makedirs, open as os_open, write as os_write, close as os_close, O_CREAT, O_APPEND, O_WRONLY
+from os.path import dirname, isfile
 from select import select
 from socket import create_connection, SHUT_RDWR
+from ssl import HAS_SNI, SSLError
 from struct import pack as struct_pack, unpack as struct_unpack, unpack_from as struct_unpack_from
 
-from ..meta import version
-from .compat import hex2, secure_socket
+from .constants import DEFAULT_PORT, DEFAULT_USER_AGENT, KNOWN_HOSTS, MAGIC_PREAMBLE, \
+    SECURITY_NONE, SECURITY_TRUST_ON_FIRST_USE
+from .compat import hex2
 from .exceptions import ProtocolError
 from .packstream import Packer, Unpacker
 
 
-DEFAULT_PORT = 7687
-DEFAULT_USER_AGENT = "neo4j-python/%s" % version
-
-MAGIC_PREAMBLE = 0x6060B017
-
 # Signature bytes for each message type
 INIT = b"\x01"             # 0000 0001 // INIT <user_agent>
 RESET = b"\x0F"            # 0000 1111 // RESET
@@ -211,14 +210,18 @@ def __init__(self, sock, **config):
         user_agent = config.get("user_agent", DEFAULT_USER_AGENT)
         if isinstance(user_agent, bytes):
             user_agent = user_agent.decode("UTF-8")
+        self.user_agent = user_agent
+
+        # Pick up the server certificate, if any
+        self.der_encoded_server_certificate = config.get("der_encoded_server_certificate")
 
         def on_failure(metadata):
             raise ProtocolError("Initialisation failed")
 
         response = Response(self)
         response.on_failure = on_failure
 
-        self.append(INIT, (user_agent,), response=response)
+        self.append(INIT, (self.user_agent,), response=response)
         self.send()
         while not response.complete:
             self.fetch_next()
@@ -313,7 +316,39 @@ def close(self):
             self.closed = True
 
 
-def connect(host, port=None, **config):
+def verify_certificate(host, der_encoded_certificate):
+    base64_encoded_certificate = b64encode(der_encoded_certificate)
+    if isfile(KNOWN_HOSTS):
+        with open(KNOWN_HOSTS) as f_in:
+            for line in f_in:
+                known_host, _, known_cert = line.strip().partition(":")
+                if host == known_host:
+                    if base64_encoded_certificate == known_cert:
+                        # Certificate match
+                        return
+                    else:
+                        # Certificate mismatch
+                        print(base64_encoded_certificate)
+                        print(known_cert)
+                        raise ProtocolError("Server certificate does not match known certificate for %r; check "
+                                            "details in file %r" % (host, KNOWN_HOSTS))
+    # First use (no hosts match)
+    try:
+        makedirs(dirname(KNOWN_HOSTS))
+    except OSError:
+        pass
+    f_out = os_open(KNOWN_HOSTS, O_CREAT | O_APPEND | O_WRONLY, 0o600)  # TODO: Windows
+    if isinstance(host, bytes):
+        os_write(f_out, host)
+    else:
+        os_write(f_out, host.encode("utf-8"))
+    os_write(f_out, b":")
+    os_write(f_out, base64_encoded_certificate)
+    os_write(f_out, b"\n")
+    os_close(f_out)
+
+
+def connect(host, port=None, ssl_context=None, **config):
     """ Connect and perform a handshake and return a valid Connection object, assuming
     a protocol version can be agreed.
     """
@@ -323,10 +358,25 @@ def connect(host, port=None, **config):
     if __debug__: log_info("~~ [CONNECT] %s %d", host, port)
     s = create_connection((host, port))
 
-    # Secure the connection if so requested
-    if config.get("secure", False):
+    # Secure the connection if an SSL context has been provided
+    if ssl_context:
         if __debug__: log_info("~~ [SECURE] %s", host)
-        s = secure_socket(s, host)
+        try:
+            s = ssl_context.wrap_socket(s, server_hostname=host if HAS_SNI else None)
+        except SSLError as cause:
+            error = ProtocolError("Cannot establish secure connection; %s" % cause.args[1])
+            error.__cause__ = cause
+            raise error
+        else:
+            # Check that the server provides a certificate
+            der_encoded_server_certificate = s.getpeercert(binary_form=True)
+            if der_encoded_server_certificate is None:
+                raise ProtocolError("When using a secure socket, the server should always provide a certificate")
+            security = config.get("security", SECURITY_NONE)
+            if security == SECURITY_TRUST_ON_FIRST_USE:
+                verify_certificate(host, der_encoded_server_certificate)
+    else:
+        der_encoded_server_certificate = None
 
     # Send details of the protocol versions supported
     supported_versions = [1, 0, 0, 0]
@@ -360,4 +410,4 @@ def connect(host, port=None, **config):
         s.shutdown(SHUT_RDWR)
         s.close()
     else:
-        return Connection(s, **config)
+        return Connection(s, der_encoded_server_certificate=der_encoded_server_certificate, **config)

neo4j/v1/constants.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+# Copyright (c) 2002-2016 "Neo Technology,"
+# Network Engine for Objects in Lund AB [http://neotechnology.com]
+#
+# This file is part of Neo4j.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from os.path import expanduser, join
+
+from ..meta import version
+
+
+DEFAULT_PORT = 7687
+DEFAULT_USER_AGENT = "neo4j-python/%s" % version
+
+KNOWN_HOSTS = join(expanduser("~"), ".neo4j", "known_hosts")
+
+MAGIC_PREAMBLE = 0x6060B017
+
+SECURITY_NONE = 0
+SECURITY_TRUST_ON_FIRST_USE = 1
+SECURITY_VERIFIED = 2

neo4j/v1/session.py

@@ -29,9 +29,11 @@ class which can be used to obtain `Driver` instances that are used for
 from __future__ import division
 
 from collections import deque, namedtuple
+from ssl import SSLContext, PROTOCOL_SSLv23, OP_NO_SSLv2, CERT_REQUIRED, Purpose
 
 from .compat import integer, string, urlparse
 from .connection import connect, Response, RUN, PULL_ALL
+from .constants import SECURITY_NONE, SECURITY_VERIFIED
 from .exceptions import CypherError
 from .typesystem import hydrated
 
@@ -77,6 +79,16 @@ def __init__(self, url, **config):
         self.config = config
         self.max_pool_size = config.get("max_pool_size", DEFAULT_MAX_POOL_SIZE)
         self.session_pool = deque()
+        self.security = security = config.get("security", SECURITY_NONE)
+        if security > SECURITY_NONE:
+            ssl_context = SSLContext(PROTOCOL_SSLv23)
+            ssl_context.options |= OP_NO_SSLv2
+            if security >= SECURITY_VERIFIED:
+                ssl_context.verify_mode = CERT_REQUIRED
+            ssl_context.load_default_certs(Purpose.SERVER_AUTH)
+            self.ssl_context = ssl_context
+        else:
+            self.ssl_context = None
 
     def session(self):
         """ Create a new session based on the graph database details
@@ -417,13 +429,16 @@ class Session(object):
 
     def __init__(self, driver):
         self.driver = driver
-        self.connection = connect(driver.host, driver.port, **driver.config)
+        self.connection = connect(driver.host, driver.port, driver.ssl_context, **driver.config)
         self.transaction = None
         self.last_cursor = None
 
     def __del__(self):
-        if not self.connection.closed:
-            self.connection.close()
+        try:
+            if not self.connection.closed:
+                self.connection.close()
+        except AttributeError:
+            pass
 
     def __enter__(self):
         return self
@@ -643,6 +658,7 @@ def __eq__(self, other):
     def __ne__(self, other):
         return not self.__eq__(other)
 
+
 def record(obj):
     """ Obtain an immutable record for the given object
     (either by calling obj.__record__() or by copying out the record data)

test/test_session.py

@@ -19,15 +19,21 @@
 # limitations under the License.
 
 
+from os import remove, rename
+from os.path import isfile
 from socket import socket
 from ssl import SSLSocket
 from unittest import TestCase
 
 from mock import patch
+from neo4j.v1.constants import KNOWN_HOSTS, SECURITY_NONE, SECURITY_TRUST_ON_FIRST_USE, SECURITY_VERIFIED
 from neo4j.v1.session import GraphDatabase, CypherError, Record, record
 from neo4j.v1.typesystem import Node, Relationship, Path
 
 
+KNOWN_HOSTS_BACKUP = KNOWN_HOSTS + ".backup"
+
+
 class DriverTestCase(TestCase):
 
     def test_healthy_session_will_be_returned_to_the_pool_on_close(self):
@@ -81,17 +87,57 @@ def test_sessions_are_not_reused_if_still_in_use(self):
         session_1.close()
         assert session_1 is not session_2
 
-    def test_insecure_session_uses_insecure_socket(self):
-        driver = GraphDatabase.driver("bolt://localhost", secure=False)
+
+class SecurityTestCase(TestCase):
+
+    def setUp(self):
+        if isfile(KNOWN_HOSTS):
+            rename(KNOWN_HOSTS, KNOWN_HOSTS_BACKUP)
+
+    def tearDown(self):
+        if isfile(KNOWN_HOSTS_BACKUP):
+            rename(KNOWN_HOSTS_BACKUP, KNOWN_HOSTS)
+
+    def test_default_session_uses_security_none(self):
+        # TODO: verify this is the correct default (maybe TOFU?)
+        driver = GraphDatabase.driver("bolt://localhost")
+        assert driver.security == SECURITY_NONE
+
+    def test_insecure_session_uses_normal_socket(self):
+        driver = GraphDatabase.driver("bolt://localhost", security=SECURITY_NONE)
+        session = driver.session()
+        connection = session.connection
+        assert isinstance(connection.channel.socket, socket)
+        assert connection.der_encoded_server_certificate is None
+        session.close()
+
+    def test_tofu_session_uses_secure_socket(self):
+        driver = GraphDatabase.driver("bolt://localhost", security=SECURITY_TRUST_ON_FIRST_USE)
         session = driver.session()
-        assert isinstance(session.connection.channel.socket, socket)
+        connection = session.connection
+        assert isinstance(connection.channel.socket, SSLSocket)
+        assert connection.der_encoded_server_certificate is not None
         session.close()
 
-    def test_secure_session_uses_secure_socket(self):
-        driver = GraphDatabase.driver("bolt://localhost", secure=True)
+    def test_tofu_session_trusts_certificate_after_first_use(self):
+        driver = GraphDatabase.driver("bolt://localhost", security=SECURITY_TRUST_ON_FIRST_USE)
         session = driver.session()
-        assert isinstance(session.connection.channel.socket, SSLSocket)
+        connection = session.connection
+        certificate = connection.der_encoded_server_certificate
         session.close()
+        session = driver.session()
+        connection = session.connection
+        assert connection.der_encoded_server_certificate == certificate
+        session.close()
+
+    # TODO: Find a way to run this test
+    # def test_verified_session_uses_secure_socket(self):
+    #     driver = GraphDatabase.driver("bolt://localhost", security=SECURITY_VERIFIED)
+    #     session = driver.session()
+    #     connection = session.connection
+    #     assert isinstance(connection.channel.socket, SSLSocket)
+    #     assert connection.der_encoded_server_certificate is not None
+    #     session.close()
 
 
 class RunTestCase(TestCase):