mirror of
				https://github.com/zulip/zulip.git
				synced 2025-11-04 05:53:43 +00:00 
			
		
		
		
	tools: Allow specifying use of basic auth in ScreenshotConfig.
This commit is contained in:
		
				
					committed by
					
						
						Tim Abbott
					
				
			
			
				
	
			
			
			
						parent
						
							5ffc51382d
						
					
				
				
					commit
					d4c6195bca
				
			@@ -23,8 +23,9 @@ import django
 | 
			
		||||
django.setup()
 | 
			
		||||
 | 
			
		||||
import argparse
 | 
			
		||||
import base64
 | 
			
		||||
import subprocess
 | 
			
		||||
from typing import Any, Dict, Optional
 | 
			
		||||
from typing import Any, Dict
 | 
			
		||||
 | 
			
		||||
import requests
 | 
			
		||||
import ujson
 | 
			
		||||
@@ -98,7 +99,7 @@ def custom_headers(headers_json: str) -> Dict[str, str]:
 | 
			
		||||
            'Note: all strings must be enclosed within "" instead of \'\''.format(ve))
 | 
			
		||||
 | 
			
		||||
def send_bot_payload_message(bot: UserProfile, integration: WebhookIntegration, fixture_path: str,
 | 
			
		||||
                             extra_headers: Optional[Dict[str, str]]) -> bool:
 | 
			
		||||
                             config: ScreenshotConfig) -> bool:
 | 
			
		||||
    # Delete all messages, so new message is the only one it's message group
 | 
			
		||||
    Message.objects.filter(sender=bot).delete()
 | 
			
		||||
 | 
			
		||||
@@ -111,8 +112,13 @@ def send_bot_payload_message(bot: UserProfile, integration: WebhookIntegration,
 | 
			
		||||
        data = ujson.load(f)
 | 
			
		||||
    _, fixture_name = split_fixture_path(fixture_path)
 | 
			
		||||
    headers = get_requests_headers(integration.name, fixture_name)
 | 
			
		||||
    if extra_headers:
 | 
			
		||||
        headers.update(extra_headers)
 | 
			
		||||
    if config.custom_headers:
 | 
			
		||||
        headers.update(config.custom_headers)
 | 
			
		||||
    if config.use_basic_auth:
 | 
			
		||||
        credentials = base64.b64encode('{}:{}'.format(bot.email, bot.api_key).encode('utf8')).decode('utf8')
 | 
			
		||||
        auth = 'basic {}'.format(credentials)
 | 
			
		||||
        headers.update(dict(Authorization=auth))
 | 
			
		||||
 | 
			
		||||
    try:
 | 
			
		||||
        response = requests.post(url, json=data, headers=headers)
 | 
			
		||||
    except requests.exceptions.ConnectionError:
 | 
			
		||||
@@ -141,7 +147,7 @@ def generate_screenshot_from_config(integration_name: str, screenshot_config: Sc
 | 
			
		||||
    fixture_path, image_path = get_fixture_and_image_paths(integration, screenshot_config)
 | 
			
		||||
    bot = create_integration_bot(integration)
 | 
			
		||||
    create_integration_stream(integration, bot)
 | 
			
		||||
    message_sent = send_bot_payload_message(bot, integration, fixture_path, screenshot_config.custom_headers)
 | 
			
		||||
    message_sent = send_bot_payload_message(bot, integration, fixture_path, screenshot_config)
 | 
			
		||||
    if message_sent:
 | 
			
		||||
        capture_last_message_screenshot(bot, image_path)
 | 
			
		||||
        print(f'Screenshot captured to: {BOLDRED}{image_path}{ENDC}')
 | 
			
		||||
@@ -151,6 +157,8 @@ parser.add_argument('integration', type=str, help='Name of the integration')
 | 
			
		||||
parser.add_argument('fixture', type=str, help='Name of the fixture file to use')
 | 
			
		||||
parser.add_argument('--image-name', type=str, default='001.png', help='Name for the screenshot image')
 | 
			
		||||
parser.add_argument('--image-dir', type=str, help='Directory name where to save the screenshot image')
 | 
			
		||||
parser.add_argument('-A', '--use-basic-auth', action='store_true',
 | 
			
		||||
                    help='Add basic auth headers to the request')
 | 
			
		||||
parser.add_argument('-H', '--custom-headers',
 | 
			
		||||
                    type=custom_headers,
 | 
			
		||||
                    help='Any additional headers to be sent with the request.')
 | 
			
		||||
@@ -158,5 +166,5 @@ parser.add_argument('-H', '--custom-headers',
 | 
			
		||||
options = parser.parse_args()
 | 
			
		||||
prepare_puppeteer_run()
 | 
			
		||||
screenshot_config = ScreenshotConfig(
 | 
			
		||||
    options.fixture, options.image_name, options.image_dir, options.custom_headers)
 | 
			
		||||
    options.fixture, options.image_name, options.image_dir, options.use_basic_auth, options.custom_headers)
 | 
			
		||||
generate_screenshot_from_config(options.integration, screenshot_config)
 | 
			
		||||
 
 | 
			
		||||
@@ -200,10 +200,11 @@ def split_fixture_path(path: str) -> Tuple[str, str]:
 | 
			
		||||
# FIXME: Change to namedtuple if we drop Python3.6: No default values support on namedtuples (or dataclass)
 | 
			
		||||
class ScreenshotConfig:
 | 
			
		||||
    def __init__(self, fixture_name: str, image_name: str='001.png', image_dir: Optional[str]=None,
 | 
			
		||||
                 custom_headers: Optional[Dict[str, str]]=None):
 | 
			
		||||
                 use_basic_auth: bool=False, custom_headers: Optional[Dict[str, str]]=None):
 | 
			
		||||
        self.fixture_name = fixture_name
 | 
			
		||||
        self.image_name = image_name
 | 
			
		||||
        self.image_dir = image_dir
 | 
			
		||||
        self.use_basic_auth = use_basic_auth
 | 
			
		||||
        self.custom_headers = custom_headers
 | 
			
		||||
 | 
			
		||||
def get_fixture_and_image_paths(integration: WebhookIntegration,
 | 
			
		||||
 
 | 
			
		||||
		Reference in New Issue
	
	Block a user