1 files changed, 20 insertions, 9 deletions
diff --git a/src/test/ntor_ref.py b/src/test/ntor_ref.py
index ade468da7..12eb00742 100644..100755
--- a/src/test/ntor_ref.py
+++ b/src/test/ntor_ref.py
@@ -1,3 +1,4 @@
+#!/usr/bin/python
 # Copyright 2012-2013, The Tor Project, Inc
 # See LICENSE for licensing information
 
@@ -27,7 +28,14 @@ commands:
 """
 
 import binascii
-import curve25519
+try:
+    import curve25519
+    curve25519mod = curve25519.keys
+except ImportError:
+    curve25519 = None
+    import slownacl_curve25519
+    curve25519mod = slownacl_curve25519
+
 import hashlib
 import hmac
 import subprocess
@@ -67,17 +75,17 @@ T_VERIFY = PROTOID + ":verify"
 def H_mac(msg): return H(msg, tweak=T_MAC)
 def H_verify(msg): return H(msg, tweak=T_VERIFY)
 
-class PrivateKey(curve25519.keys.Private):
-    """As curve25519.keys.Private, but doesn't regenerate its public key
+class PrivateKey(curve25519mod.Private):
+    """As curve25519mod.Private, but doesn't regenerate its public key
        every time you ask for it.
     """
     def __init__(self):
-        curve25519.keys.Private.__init__(self)
+        curve25519mod.Private.__init__(self)
         self._memo_public = None
 
     def get_public(self):
         if self._memo_public is None:
-            self._memo_public = curve25519.keys.Private.get_public(self)
+            self._memo_public = curve25519mod.Private.get_public(self)
 
         return self._memo_public
 
@@ -177,7 +185,7 @@ def server(seckey_b, my_node_id, message, keyBytes=72):
     badness = (keyid(seckey_b.get_public()) !=
                message[NODE_ID_LENGTH:NODE_ID_LENGTH+H_LENGTH])
 
-    pubkey_X = curve25519.keys.Public(message[NODE_ID_LENGTH+H_LENGTH:])
+    pubkey_X = curve25519mod.Public(message[NODE_ID_LENGTH+H_LENGTH:])
     seckey_y = PrivateKey()
     pubkey_Y = seckey_y.get_public()
     pubkey_B = seckey_b.get_public()
@@ -240,7 +248,7 @@ def client_part2(seckey_x, msg, node_id, pubkey_B, keyBytes=72):
     """
     assert len(msg) == G_LENGTH + H_LENGTH
 
-    pubkey_Y = curve25519.keys.Public(msg[:G_LENGTH])
+    pubkey_Y = curve25519mod.Public(msg[:G_LENGTH])
     their_auth = msg[G_LENGTH:]
 
     pubkey_X = seckey_x.get_public()
@@ -286,6 +294,7 @@ def demo(node_id="iToldYouAboutStairs.", server_key=PrivateKey()):
     assert len(skeys) == 72
     assert len(ckeys) == 72
     assert skeys == ckeys
+    print "OK"
 
 # ======================================================================
 def timing():
@@ -368,13 +377,15 @@ def test_tor():
     assert c_keys == s_keys
     assert len(c_keys) == 90
 
-    print "We just interoperated."
+    print "OK"
 
 # ======================================================================
 
 if __name__ == '__main__':
     import sys
-    if sys.argv[1] == 'gen_kdf_vectors':
+    if len(sys.argv) < 2:
+        print __doc__
+    elif sys.argv[1] == 'gen_kdf_vectors':
         kdf_vectors()
     elif sys.argv[1] == 'timing':
         timing()