Set url-agent to Lorry if we get 403 HTTP status

If we get an HTTPError 403 then set the user-agent so it is clear that
the request is coming from Lorry. This is used to get around situations where
the url has blocked python user-agents.

1 files changed, 19 insertions, 4 deletions

diff --git a/lorry b/lorry
index cadc0f2..c3c6208 100755
--- a/lorry
+++ b/lorry
@@ -33,6 +33,7 @@ import ftplib
 import re
 import subprocess
 import tempfile
+import contextlib
 
 import yaml
 
@@ -809,7 +810,7 @@ class Lorry(cliapp.Application):
                 self.progress('.. attempting to fetch %s' % basename)
                 try:
                     with open(file_dest, 'wb') as raw_file, \
-                         urllib.request.urlopen(url) as urlfile:
+                        self.urlopen(url) as urlfile:
                         raw_file.write(urlfile.read())
                         try:
                             # HTTP dates use (one of) the email date formats
@@ -857,8 +858,8 @@ class Lorry(cliapp.Application):
         if file_missing_or_empty(archive_dest):
             self.progress('.. attempting to fetch.')
             try:
-                with open(archive_dest, 'wb') as archive_file:
-                    urlfile = urllib.request.urlopen(spec['url'])
+                with open(archive_dest, 'wb') as archive_file, \
+                    self.urlopen(spec['url']) as urlfile:
                     archive_file.write(urlfile.read())
                     try:
                         # HTTP dates use (one of) the email date formats
@@ -867,7 +868,6 @@ class Lorry(cliapp.Application):
                                 urlfile.info()['Last-Modified']))
                     except (KeyError, ValueError, TypeError):
                         url_date = None
-                    urlfile.close()
                 if url_date:
                     os.utime(archive_dest, (url_date, url_date))
             except Exception:
@@ -984,6 +984,21 @@ class Lorry(cliapp.Application):
             os.unlink(marks_temp_name)
             raise
 
+    @contextlib.contextmanager
+    def urlopen(self, url):
+        try:
+            req = urllib.request.Request(url)
+            with urllib.request.urlopen(req) as urlfile:
+                yield urlfile
+        except urllib.error.HTTPError as e:
+            if e.getcode() == 403:
+                newreq = urllib.request.Request(url)
+                newreq.add_header('User-Agent',
+                    'Lorry/%s (https://gitlab.com/CodethinkLabs/lorry/lorry)' % __version__)
+                with urllib.request.urlopen(newreq) as newurlfile:
+                    yield newurlfile
+            else:
+                raise
 
 if __name__ == '__main__':
     Lorry(version=__version__).run()