From 384626cff968366d043719f702e82571a17ec1e8 Mon Sep 17 00:00:00 2001
From: ttgump <syrttgump@gmail.com>
Date: Fri, 13 Sep 2013 19:14:32 -0400
Subject: [PATCH] test_mix_and_match_attack and test_indefinite_freeze_attack
 refactory fix

---
 .../test_indefinite_freeze_attack.py          | 40 +++++++++----------
 .../integration/test_mix_and_match_attack.py  | 20 +++++-----
 2 files changed, 30 insertions(+), 30 deletions(-)

diff --git a/tests/integration/test_indefinite_freeze_attack.py b/tests/integration/test_indefinite_freeze_attack.py
index b422fae5..2b88e00c 100755
--- a/tests/integration/test_indefinite_freeze_attack.py
+++ b/tests/integration/test_indefinite_freeze_attack.py
@@ -28,7 +28,7 @@
 
 import tuf
 import tuf.formats
-import tuf.interposition.urllib_tuf as urllib_tuf
+import tuf.interposition
 import tuf.repo.signerlib as signerlib
 import tuf.tests.util_test_tools as util_test_tools
 
@@ -61,9 +61,9 @@ def _remake_timestamp(metadata_dir, keyids):
 
 
 
-def _download(url, filename, tuf=False):
-  if tuf:
-    urllib_tuf.urlretrieve(url, filename)
+def _download(url, filename, using_tuf=False):
+  if using_tuf:
+    tuf.interposition.urllib_tuf.urlretrieve(url, filename)
     
   else:
     urllib.urlretrieve(url, filename)
@@ -117,26 +117,24 @@ def test_indefinite_freeze_attack(TUF=False):
       _remake_timestamp(metadata_dir, keyids)
 
 
-    # Client performs initial download.
+    # Client performs initial download. If the computer is slow, it may
+    # take longer time than expiration time. In this case you will see
+    # an ExpiredMetadataError.
     try:
-      _download(url=url_to_repo, filename=downloaded_file, tuf=TUF)
-    except tuf.ExpiredMetadataError:
-      msg = ('Metadata has expired too soon, extend expiration period. '+
-             'Current expiration is set to: '+repr(EXPIRATION)+' second(s).')
-      sys.exit(msg)
-
-    # Expire timestamp.
-    time.sleep(EXPIRATION)
-
-    # Try downloading again, this should raise an error.
-    try:
-      _download(url=url_to_repo, filename=downloaded_file, tuf=TUF)
-    except tuf.ExpiredMetadataError, error:
-      pass
+      _download(url=url_to_repo, filename=downloaded_file, using_tuf=TUF)
+    except:
+      print ('Initial download failed! It may be because your machine is busy. Try it later.')
     else:
-      raise IndefiniteFreezeAttackAlert(ERROR_MSG)
-
+      # Expire timestamp.
+      time.sleep(EXPIRATION)
 
+      # Try downloading again, this should raise an error.
+      try:
+	_download(url=url_to_repo, filename=downloaded_file, using_tuf=TUF)
+      except tuf.ExpiredMetadataError, error:
+	print('Caught an expiration error!')
+      else:
+	raise IndefiniteFreezeAttackAlert(ERROR_MSG)
   finally:
     util_test_tools.cleanup(root_repo, server_proc)
 
diff --git a/tests/integration/test_mix_and_match_attack.py b/tests/integration/test_mix_and_match_attack.py
index da8e170c..88cbaeb8 100755
--- a/tests/integration/test_mix_and_match_attack.py
+++ b/tests/integration/test_mix_and_match_attack.py
@@ -40,7 +40,7 @@
 import tempfile
 
 import tuf
-import tuf.interposition.urllib_tuf as urllib_tuf
+import tuf.interposition
 import tuf.tests.util_test_tools as util_test_tools
 
 
@@ -48,9 +48,9 @@ class MixAndMatchAttackAlert(Exception):
   pass
 
 
-def _download(url, filename, tuf=False):
-  if tuf:
-    urllib_tuf.urlretrieve(url, filename)
+def _download(url, filename, using_tuf=False):
+  if using_tuf:
+    tuf.interposition.urllib_tuf.urlretrieve(url, filename)
     
   else:
     urllib.urlretrieve(url, filename)
@@ -126,7 +126,7 @@ def test_mix_and_match_attack(TUF=False):
 
 
     # Client's initial download.
-    _download(url=url_to_file, filename=downloaded_file, tuf=TUF)
+    _download(url=url_to_file, filename=downloaded_file, using_tuf=TUF)
 
     # Stage 2
     # -------
@@ -139,7 +139,7 @@ def test_mix_and_match_attack(TUF=False):
       util_test_tools.tuf_refresh_repo(root_repo, keyids)
 
     # Client downloads the patched file.
-    _download(url=url_to_file, filename=downloaded_file, tuf=TUF)
+    _download(url=url_to_file, filename=downloaded_file, using_tuf=TUF)
 
     downloaded_content = util_test_tools.read_file_content(downloaded_file)
 
@@ -163,9 +163,11 @@ def test_mix_and_match_attack(TUF=False):
 
     # Client tries to downloads the newly patched file.
     try:
-      _download(url=url_to_file, filename=downloaded_file, tuf=TUF)
-    except tuf.MetadataNotAvailableError:
-      pass
+      _download(url=url_to_file, filename=downloaded_file, using_tuf=TUF)
+    except tuf.NoWorkingMirrorError as errors:
+      for mirror_url, mirror_error in errors.mirror_errors.iteritems():
+	if type(mirror_error) == tuf.BadHashError:
+	  print 'Catched a Bad Hash Error!'
 
     # Check whether the attack succeeded by inspecting the content of the
     # update.  The update should contain 'Test NOT A'.