From a4549b44d835afa3bbf50f550e50f07f181944ac Mon Sep 17 00:00:00 2001 From: Maxime Courtet Date: Tue, 30 Apr 2019 15:46:42 +0200 Subject: [PATCH 1/3] Update generate_bytes_img_functions From now on supports credentials={} --- screamshot/generate_bytes_img_functions.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/screamshot/generate_bytes_img_functions.py b/screamshot/generate_bytes_img_functions.py index 84ba6b0..f8a395e 100644 --- a/screamshot/generate_bytes_img_functions.py +++ b/screamshot/generate_bytes_img_functions.py @@ -40,11 +40,12 @@ def _parse_parameters(**kwargs) -> dict: 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_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, From e81baa66f964300d383bdc91c10c85dcda8fd366 Mon Sep 17 00:00:00 2001 From: Maxime Courtet Date: Tue, 30 Apr 2019 15:47:18 +0200 Subject: [PATCH 2/3] Update screamshot script file Crendentials can be specified --- screamshot/screamshot_script.py | 34 ++++++++++++++++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/screamshot/screamshot_script.py b/screamshot/screamshot_script.py index e798525..d690e58 100644 --- a/screamshot/screamshot_script.py +++ b/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,6 +69,25 @@ 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[^<]+)\:(?P[^<]+)', 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)) @@ -64,7 +95,8 @@ def main(): 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()) From 2ac2193e0322452f594dccaed2afc937c1c8eb1c Mon Sep 17 00:00:00 2001 From: Maxime Courtet Date: Tue, 30 Apr 2019 16:13:38 +0200 Subject: [PATCH 3/3] Update tests --- tests/test_screamshot_script.py | 48 ++++++++++++++++++++++++++++++++- 1 file changed, 47 insertions(+), 1 deletion(-) diff --git a/tests/test_screamshot_script.py b/tests/test_screamshot_script.py index e7de2a6..7c582ca 100644 --- a/tests/test_screamshot_script.py +++ b/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()