Security TOFU, PersonalCertificateStore

Author: technige

neo4j/v1/connection.py

@@ -33,7 +33,7 @@
 from struct import pack as struct_pack, unpack as struct_unpack, unpack_from as struct_unpack_from
 
 from .constants import DEFAULT_PORT, DEFAULT_USER_AGENT, KNOWN_HOSTS, MAGIC_PREAMBLE, \
-    SECURITY_NONE, SECURITY_TRUST_ON_FIRST_USE
+    SECURITY_DEFAULT, SECURITY_TRUST_ON_FIRST_USE
 from .compat import hex2
 from .exceptions import ProtocolError
 from .packstream import Packer, Unpacker
@@ -316,36 +316,51 @@ def close(self):
             self.closed = True
 
 
-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)
+class CertificateStore(object):
+
+    def match_or_trust(self, host, der_encoded_certificate):
+        """ Check whether the supplied certificate matches that stored for the
+        specified host. If it does, return ``True``, if it doesn't, return
+        ``False``. If no entry for that host is found, add it to the store
+        and return ``True``.
+
+        :arg host:
+        :arg der_encoded_certificate:
+        :return:
+        """
+        raise NotImplementedError()
+
+
+class PersonalCertificateStore(CertificateStore):
+
+    def __init__(self, path=None):
+        self.path = path or KNOWN_HOSTS
+
+    def match_or_trust(self, host, der_encoded_certificate):
+        base64_encoded_certificate = b64encode(der_encoded_certificate)
+        if isfile(self.path):
+            with open(self.path) as f_in:
+                for line in f_in:
+                    known_host, _, known_cert = line.strip().partition(":")
+                    if host == known_host:
+                        print("Received: %s" % base64_encoded_certificate)
+                        print("Known: %s" % known_cert)
+                        return base64_encoded_certificate == known_cert
+        # First use (no hosts match)
+        try:
+            makedirs(dirname(self.path))
+        except OSError:
+            pass
+        f_out = os_open(self.path, 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)
+        return True
 
 
 def connect(host, port=None, ssl_context=None, **config):
@@ -372,9 +387,12 @@ def connect(host, port=None, ssl_context=None, **config):
             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)
+            security = config.get("security", SECURITY_DEFAULT)
             if security == SECURITY_TRUST_ON_FIRST_USE:
-                verify_certificate(host, der_encoded_server_certificate)
+                store = PersonalCertificateStore()
+                if not store.match_or_trust(host, der_encoded_server_certificate):
+                    raise ProtocolError("Server certificate does not match known certificate for %r; check "
+                                        "details in file %r" % (host, KNOWN_HOSTS))
     else:
         der_encoded_server_certificate = None
 

neo4j/v1/constants.py

@@ -34,3 +34,5 @@
 SECURITY_NONE = 0
 SECURITY_TRUST_ON_FIRST_USE = 1
 SECURITY_VERIFIED = 2
+
+SECURITY_DEFAULT = SECURITY_TRUST_ON_FIRST_USE

neo4j/v1/session.py

@@ -33,7 +33,7 @@ class which can be used to obtain `Driver` instances that are used for
 
 from .compat import integer, string, urlparse
 from .connection import connect, Response, RUN, PULL_ALL
-from .constants import SECURITY_NONE, SECURITY_VERIFIED
+from .constants import SECURITY_NONE, SECURITY_VERIFIED, SECURITY_DEFAULT
 from .exceptions import CypherError, ResultError
 from .typesystem import hydrated
 
@@ -79,7 +79,7 @@ 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)
+        self.security = security = config.get("security", SECURITY_DEFAULT)
         if security > SECURITY_NONE:
             ssl_context = SSLContext(PROTOCOL_SSLv23)
             ssl_context.options |= OP_NO_SSLv2

test/test_session.py

@@ -19,23 +19,19 @@
 # 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.constants import SECURITY_NONE, SECURITY_TRUST_ON_FIRST_USE
 from neo4j.v1.exceptions import CypherError, ResultError
 from neo4j.v1.session import GraphDatabase, Record, record
 from neo4j.v1.typesystem import Node, Relationship, Path
 
+from test.util import ServerTestCase
 
-KNOWN_HOSTS_BACKUP = KNOWN_HOSTS + ".backup"
 
-
-class DriverTestCase(TestCase):
+class DriverTestCase(ServerTestCase):
 
     def test_healthy_session_will_be_returned_to_the_pool_on_close(self):
         driver = GraphDatabase.driver("bolt://localhost")
@@ -89,20 +85,11 @@ def test_sessions_are_not_reused_if_still_in_use(self):
         assert session_1 is not session_2
 
 
-class SecurityTestCase(TestCase):
-
-    def setUp(self):
-        if isfile(KNOWN_HOSTS):
-            rename(KNOWN_HOSTS, KNOWN_HOSTS_BACKUP)
+class SecurityTestCase(ServerTestCase):
 
-    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?)
+    def test_default_session_uses_tofu(self):
         driver = GraphDatabase.driver("bolt://localhost")
-        assert driver.security == SECURITY_NONE
+        assert driver.security == SECURITY_TRUST_ON_FIRST_USE
 
     def test_insecure_session_uses_normal_socket(self):
         driver = GraphDatabase.driver("bolt://localhost", security=SECURITY_NONE)
@@ -141,7 +128,7 @@ def test_tofu_session_trusts_certificate_after_first_use(self):
     #     session.close()
 
 
-class RunTestCase(TestCase):
+class RunTestCase(ServerTestCase):
 
     def test_can_run_simple_statement(self):
         session = GraphDatabase.driver("bolt://localhost").session()
@@ -363,7 +350,7 @@ def test_can_obtain_notification_info(self):
             assert position.column == 1
 
 
-class ResetTestCase(TestCase):
+class ResetTestCase(ServerTestCase):
 
     def test_automatic_reset_after_failure(self):
         with GraphDatabase.driver("bolt://localhost").session() as session:
@@ -387,7 +374,7 @@ def test_defunct(self):
             assert session.connection.closed
 
 
-class RecordTestCase(TestCase):
+class RecordTestCase(ServerTestCase):
     def test_record_equality(self):
         record1 = Record(["name", "empire"], ["Nigel", "The British Empire"])
         record2 = Record(["name", "empire"], ["Nigel", "The British Empire"])
@@ -461,7 +448,8 @@ def test_record_repr(self):
         assert repr(a_record) == "<Record name='Nigel' empire='The British Empire'>"
 
 
-class TransactionTestCase(TestCase):
+class TransactionTestCase(ServerTestCase):
+
     def test_can_commit_transaction(self):
         with GraphDatabase.driver("bolt://localhost").session() as session:
             tx = session.begin_transaction()

test/util.py

@@ -20,8 +20,15 @@
 
 
 import functools
+from os import rename
+from os.path import isfile
+from unittest import TestCase
 
 from neo4j.util import Watcher
+from neo4j.v1.constants import KNOWN_HOSTS
+
+
+KNOWN_HOSTS_BACKUP = KNOWN_HOSTS + ".backup"
 
 
 def watch(f):
@@ -39,3 +46,19 @@ def wrapper(*args, **kwargs):
         f(*args, **kwargs)
         watcher.stop()
     return wrapper
+
+
+class ServerTestCase(TestCase):
+    """ Base class for test cases that use a remote server.
+    """
+
+    known_hosts = KNOWN_HOSTS
+    known_hosts_backup = known_hosts + ".backup"
+
+    def setUp(self):
+        if isfile(self.known_hosts):
+            rename(self.known_hosts, self.known_hosts_backup)
+
+    def tearDown(self):
+        if isfile(self.known_hosts_backup):
+            rename(self.known_hosts_backup, self.known_hosts)