Merge 94027c4 into 2d816a1

screamshot/generate_bytes_img_functions.py

@@ -38,13 +38,14 @@ def _parse_parameters(**kwargs) -> dict:
         wait_for = kwargs.pop("wait_for")
 
     credentials = {}
-    if "credentials" in kwargs:
-        credentials_data = kwargs.pop("credentials")
-        if "username" in credentials_data and "password" in credentials_data:
-            credentials["login"] = True
-        if "token_in_header" in credentials_data:
-            credentials["token_in_header"] = credentials_data.pop("token_in_header")
-        credentials.update({"credentials_data": credentials_data})
+    if 'credentials' in kwargs:
+        credentials_data = kwargs.pop('credentials')
+        if credentials_data:
+            if 'username' in credentials_data and 'password' in credentials_data:
+                credentials['login'] = True
+            if 'token_in_header' in credentials_data:
+                credentials['token_in_header'] = credentials_data.pop('token_in_header')
+            credentials.update({'credentials_data': credentials_data})
 
     return {
         "arg_viewport": arg_viewport,

screamshot/screamshot_script.py

@@ -2,11 +2,16 @@
 Take a screenshot
 """
 from argparse import ArgumentParser
+from logging import getLogger
+from re import search
 
 from screamshot import generate_bytes_img
 from screamshot.utils import to_sync, open_browser, close_browser
 
 
+logger = getLogger('screamshot')
+
+
 def main():
     parser = ArgumentParser(description=__doc__)
 
@@ -31,6 +36,13 @@ def main():
                         default="load",
                         help="How long do you want to wait for the page to be loaded")
 
+    # Credentials group
+    credentials_group = parser.add_argument_group(title="Credentials (optional)")
+    credentials_group.add_argument('--username', help='The username to use')
+    credentials_group.add_argument('--password', help='The password to use')
+    credentials_group.add_argument('--token', help='The header line to add. \
+        Must be like the following expression: key:token')
+
     # CSS3 selectors group
     selector_group = parser.add_argument_group(title="CSS3 selectors (optional)",
                                                description="Using quotes is recommended for \
@@ -57,14 +69,34 @@ def main():
 
     args = parser.parse_args()
 
+    credentials = None
+    if args.username and not args.password:
+        logger.error('A password must be specified')
+        exit(1)
+    elif not args.username and args.password:
+        logger.error('A username must be specified')
+        exit(1)
+    elif args.username and args.password:
+        credentials = {'username': args.username, 'password': args.password}
+    elif args.token:
+        regex_token = search(r'(?P<key>[^<]+)\:(?P<token>[^<]+)', args.token)
+        try:
+            key = regex_token.group('key')
+            token = regex_token.group('token')
+            credentials = {key: token, 'token_in_header': True}
+        except AttributeError as _:
+            logger.error('Bad token argument, please read the documentation')
+            exit(1)
+
     if not args.no_browser:
         to_sync(
             open_browser(args.headless, launch_args=args.no_sandbox))
 
     to_sync(
         generate_bytes_img(args.url, path=args.path, width=args.width, height=args.height,
                            full_page=args.fullpage, selector=args.selector,
-                           wait_for=args.wait_for, wait_until=args.wait_until))
+                           wait_for=args.wait_for, wait_until=args.wait_until,
+                           credentials=credentials))
 
     if not args.no_browser and not args.no_close:
         to_sync(close_browser())

tests/test_screamshot_script.py

@@ -2,10 +2,14 @@
 Tests screamshot script
 """
 from unittest import TestCase, main
-from subprocess import run
+from subprocess import run, PIPE
 from os import remove
 
 
+TOKEN = 'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJ1c2VyIjoibWFraW5hIn0.\
+jUTxi6c2-o3nHJ6Bq7zRXFoKixUyYetgPX3cToOayiA'
+
+
 class TestScreamshotScript(TestCase):
     """
     Test class
@@ -21,6 +25,48 @@ def test_simple_screamshot(self):
         open('test_simple_screamshot.png', 'r')
         remove('test_simple_screamshot.png')
 
+    def test_screamshot_username_without_password(self):
+        res = run(['python3', 'screamshot/screamshot_script.py', '--no-sandbox',
+                   '--username=makina', 'http://localhost:5000/protected_other',
+                   'test_simple_screamshot.png'], stderr=PIPE)
+        self.assertEqual(res.stderr.decode('utf-8'), 'A password must be specified\n')
+
+    def test_screamshot_password_without_username(self):
+        res = run(['python3', 'screamshot/screamshot_script.py', '--no-sandbox',
+                   '--password=makina', 'http://localhost:5000/protected_other',
+                   'test_simple_screamshot.png'], stderr=PIPE)
+        self.assertEqual(res.stderr.decode('utf-8'), 'A username must be specified\n')
+
+    def test_screamshot_with_username_and_password(self):
+        """
+        Takes a screenshot thanks to the script and checks if the image has been saved
+        """
+        with self.assertRaises(FileNotFoundError):
+            open('test_simple_screamshot.png', 'r')
+        run(['python3', 'screamshot/screamshot_script.py', '--no-sandbox', '--username=makina',
+             '--password=makina', 'http://localhost:5000/other.html', 'test_simple_screamshot.png'])
+        open('test_simple_screamshot.png', 'r')
+        remove('test_simple_screamshot.png')
+
+    def test_screamshot_with_bad_token(self):
+        res = run(['python3', 'screamshot/screamshot_script.py', '--no-sandbox',
+                   '--token=x!x', 'http://localhost:5000/protected_other',
+                   'test_simple_screamshot.png'], stderr=PIPE)
+        self.assertEqual(res.stderr.decode('utf-8'),
+                         'Bad token argument, please read the documentation\n')
+
+    def test_screamshot_with_token(self):
+        """
+        Takes a screenshot thanks to the script and checks if the image has been saved
+        """
+        with self.assertRaises(FileNotFoundError):
+            open('test_simple_screamshot.png', 'r')
+        run(['python3', 'screamshot/screamshot_script.py', '--no-sandbox',
+             '--token=token:{0}'.format(TOKEN), 'http://localhost:5000/other.html',
+             'test_simple_screamshot.png'])
+        open('test_simple_screamshot.png', 'r')
+        remove('test_simple_screamshot.png')
+
 
 if __name__ == '__main__':
     main()