1 files changed, 61 insertions, 23 deletions

diff --git a/setup.py b/setup.py
index 681b9031..c1a6d642 100644
--- a/setup.py
+++ b/setup.py
@@ -10,12 +10,20 @@
 # implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from __future__ import absolute_import, division, print_function
+
 import os
 import sys
 from distutils.command.build import build
+import subprocess
+
+import pkg_resources
+
+from setuptools import find_packages, setup
+from setuptools.command.install import install
+from setuptools.command.test import test
 
-from setuptools import setup, find_packages
-from setuptools.command.test import test as TestCommand
 
 base_dir = os.path.dirname(__file__)
 
@@ -26,18 +34,44 @@ with open(os.path.join(base_dir, "cryptography", "__about__.py")) as f:
 
 CFFI_DEPENDENCY = "cffi>=0.8"
 SIX_DEPENDENCY = "six>=1.4.1"
+VECTORS_DEPENDENCY = "cryptography_vectors=={0}".format(about['__version__'])
 
 requirements = [
     CFFI_DEPENDENCY,
     SIX_DEPENDENCY
 ]
 
+# If you add a new dep here you probably need to add it in the tox.ini as well
 test_requirements = [
     "pytest",
     "pretend",
-    "iso8601"
+    "iso8601",
 ]
 
+# If there's no vectors locally that probably means we are in a tarball and
+# need to go and get the matching vectors package from PyPi
+if not os.path.exists(os.path.join(base_dir, "vectors/setup.py")):
+    test_requirements.append(VECTORS_DEPENDENCY)
+
+
+def get_ext_modules():
+    from cryptography.hazmat.bindings.commoncrypto.binding import (
+        Binding as CommonCryptoBinding
+    )
+    from cryptography.hazmat.bindings.openssl.binding import (
+        Binding as OpenSSLBinding
+    )
+    from cryptography.hazmat.primitives import constant_time, padding
+
+    ext_modules = [
+        OpenSSLBinding().ffi.verifier.get_extension(),
+        constant_time._ffi.verifier.get_extension(),
+        padding._ffi.verifier.get_extension()
+    ]
+    if CommonCryptoBinding.is_available():
+        ext_modules.append(CommonCryptoBinding().ffi.verifier.get_extension())
+    return ext_modules
+
 
 class CFFIBuild(build):
     """
@@ -50,33 +84,35 @@ class CFFIBuild(build):
     """
 
     def finalize_options(self):
-        from cryptography.hazmat.bindings.commoncrypto.binding import (
-            Binding as CommonCryptoBinding
-        )
-        from cryptography.hazmat.bindings.openssl.binding import (
-            Binding as OpenSSLBinding
-        )
-        from cryptography.hazmat.primitives import constant_time, padding
-
-        self.distribution.ext_modules = [
-            OpenSSLBinding().ffi.verifier.get_extension(),
-            constant_time._ffi.verifier.get_extension(),
-            padding._ffi.verifier.get_extension()
-        ]
-        if CommonCryptoBinding.is_available():
-            self.distribution.ext_modules.append(
-                CommonCryptoBinding().ffi.verifier.get_extension()
-            )
-
+        self.distribution.ext_modules = get_ext_modules()
         build.finalize_options(self)
 
 
-class PyTest(TestCommand):
+class CFFIInstall(install):
+    """
+    As a consequence of CFFIBuild and it's late addition of ext_modules, we
+    need the equivalent for the ``install`` command to install into platlib
+    install-dir rather than purelib.
+    """
+
+    def finalize_options(self):
+        self.distribution.ext_modules = get_ext_modules()
+        install.finalize_options(self)
+
+
+class PyTest(test):
     def finalize_options(self):
-        TestCommand.finalize_options(self)
+        test.finalize_options(self)
         self.test_args = []
         self.test_suite = True
 
+        # This means there's a vectors/ folder with the package in here.
+        # cd into it, install the vectors package and then refresh sys.path
+        if VECTORS_DEPENDENCY not in test_requirements:
+            subprocess.Popen([sys.executable, "setup.py", "install"],
+                             cwd="vectors").communicate()
+            pkg_resources.get_distribution("cryptography_vectors").activate()
+
     def run_tests(self):
         # Import here because in module scope the eggs are not loaded.
         import pytest
@@ -116,6 +152,7 @@ setup(
         "Programming Language :: Python :: 3",
         "Programming Language :: Python :: 3.2",
         "Programming Language :: Python :: 3.3",
+        "Programming Language :: Python :: 3.4",
         "Programming Language :: Python :: Implementation :: CPython",
         "Programming Language :: Python :: Implementation :: PyPy",
         "Topic :: Security :: Cryptography",
@@ -132,6 +169,7 @@ setup(
     ext_package="cryptography",
     cmdclass={
         "build": CFFIBuild,
+        "install": CFFIInstall,
         "test": PyTest,
     }
 )