From c5c5ffcfbd2522f8fba11ce5b7ca09b9b891e9ce Mon Sep 17 00:00:00 2001
From: Nazarii Hnydyn <nazariig@mellanox.com>
Date: Sun, 1 Mar 2020 15:22:15 +0200
Subject: [PATCH] [fwutil]: Set default socket timeout for FW download to 30
 sec. (#821)

Signed-off-by: Nazarii Hnydyn <nazariig@mellanox.com>
---
 fwutil/lib.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/fwutil/lib.py b/fwutil/lib.py
index 3f6ecba4f303..60377c283a75 100755
--- a/fwutil/lib.py
+++ b/fwutil/lib.py
@@ -9,6 +9,7 @@
     import click
     import os
     import json
+    import socket
     import urllib
     import subprocess
     import sonic_device_util
@@ -45,7 +46,8 @@ class URL(object):
     PB_INFO_SEPARATOR = " | "
     PB_FULL_TERMINAL_WIDTH = 0
 
-    TMP_PATH = "/tmp"
+    DOWNLOAD_TIMEOUT = 30
+    DOWNLOAD_PATH_TEMPLATE = "/tmp/{}"
 
     def __init__(self, url):
         self.__url = url
@@ -115,13 +117,21 @@ def retrieve(self):
         if not extension:
             raise RuntimeError("Filename is malformed: did not find an extension")
 
+        default_timeout = socket.getdefaulttimeout()
+        socket.setdefaulttimeout(self.DOWNLOAD_TIMEOUT)
+
         try:
             filename, headers = urllib.urlretrieve(
                 self.__url,
-                "{}/{}".format(self.TMP_PATH, basename),
+                self.DOWNLOAD_PATH_TEMPLATE.format(basename),
                 self.__reporthook
             )
+        except:
+            if os.path.exists(self.DOWNLOAD_PATH_TEMPLATE.format(basename)):
+                os.remove(self.DOWNLOAD_PATH_TEMPLATE.format(basename))
+            raise
         finally:
+            socket.setdefaulttimeout(default_timeout)
             self.__pb_reset()
 
         return filename, headers