From d4c6195bcafc42df984e2e5141a96e7fd788518f Mon Sep 17 00:00:00 2001
From: Puneeth Chaganti <punchagan@muse-amuse.in>
Date: Fri, 1 May 2020 11:58:25 +0530
Subject: [PATCH] tools: Allow specifying use of basic auth in
 ScreenshotConfig.

---
 tools/generate-integration-docs-screenshot | 20 ++++++++++++++------
 zerver/lib/integrations.py                 |  3 ++-
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/tools/generate-integration-docs-screenshot b/tools/generate-integration-docs-screenshot
index c865e3992e..0ecd55393b 100755
--- a/tools/generate-integration-docs-screenshot
+++ b/tools/generate-integration-docs-screenshot
@@ -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)
diff --git a/zerver/lib/integrations.py b/zerver/lib/integrations.py
index 362ca58f43..93ffbe17f1 100644
--- a/zerver/lib/integrations.py
+++ b/zerver/lib/integrations.py
@@ -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,