Add proxy support for files.download() (pyinfra-dev#998)

* Add proxy support for files.download() Proxy support should work with curl and wget Help needed with tests * Fix too long doc line * Add `files.download` proxy tests --------- Co-authored-by: Nick Mills-Barrett <[email protected]>
sysadmin75 · Jan 13, 2024 · f0297cd · f0297cd
1 parent 8e2c991
commit f0297cd
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 0 deletions.
diff --git a/pyinfra/operations/files.py b/pyinfra/operations/files.py
@@ -69,6 +69,7 @@ def download(
     md5sum=None,
     headers=None,
     insecure=False,
+    proxy=None,
 ):
     """
     Download files from remote locations using ``curl`` or ``wget``.
@@ -85,6 +86,7 @@ def download(
     + md5sum: md5 hash to checksum the downloaded file against
     + headers: optional dictionary of headers to set for the HTTP request
     + insecure: disable SSL verification for the HTTP request
+    + proxy: simple HTTP proxy through which we can download files, form `http://<yourproxy>:<port>`
 
     **Example:**
 
@@ -140,7 +142,14 @@ def download(
         temp_file = state.get_temp_filename(dest)
 
         curl_args = ["-sSLf"]
+
         wget_args = ["-q"]
+
+        if proxy:
+            curl_args.append(f"--proxy {proxy}")
+            wget_args.append("-e use_proxy=yes")
+            wget_args.append(f"-e http_proxy={proxy}")
+
         if insecure:
             curl_args.append("--insecure")
             wget_args.append("--no-check-certificate")

diff --git a/tests/operations/files.download/download_proxy_curl.json b/tests/operations/files.download/download_proxy_curl.json
@@ -0,0 +1,20 @@
+{
+    "args": ["http://myfile"],
+    "kwargs": {
+        "dest": "/home/myfile",
+        "proxy": "socks5://someproxy"
+    },
+    "facts": {
+        "server.Date": "datetime:2015-01-01T00:00:00",
+        "files.File": {
+            "path=/home/myfile": null
+        },
+        "server.Which": {
+            "command=curl": "yes"
+        }
+    },
+    "commands": [
+        "curl -sSLf --proxy socks5://someproxy http://myfile -o _tempfile_",
+        "mv _tempfile_ /home/myfile"
+    ]
+}
diff --git a/tests/operations/files.download/download_proxy_wget.json b/tests/operations/files.download/download_proxy_wget.json
@@ -0,0 +1,21 @@
+{
+    "args": ["http://myfile"],
+    "kwargs": {
+        "dest": "/home/myfile",
+        "proxy": "socks5://someproxy"
+    },
+    "facts": {
+        "server.Date": "datetime:2015-01-01T00:00:00",
+        "files.File": {
+            "path=/home/myfile": null
+        },
+        "server.Which": {
+            "command=curl": null,
+            "command=wget": "yes"
+        }
+    },
+    "commands": [
+        "wget -q -e use_proxy=yes -e http_proxy=socks5://someproxy http://myfile -O _tempfile_ || ( rm -f _tempfile_ ; exit 1 )",
+        "mv _tempfile_ /home/myfile"
+    ]
+}