From 86bfe129013417e72c93c8707399d13219e8639b Mon Sep 17 00:00:00 2001
From: Wes Mason <wes@1stvamp.org>
Date: Mon, 11 Dec 2023 16:11:08 +0000
Subject: [PATCH] Use HTTP basic auth to run queries as wikiuser

---
 wikipedia/track.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/wikipedia/track.py b/wikipedia/track.py
index e6e541d1..8b8daa2a 100644
--- a/wikipedia/track.py
+++ b/wikipedia/track.py
@@ -3,6 +3,7 @@
 import random
 import re
 import uuid
+from base64 import b64encode
 from os import getcwd
 from os.path import dirname
 from typing import Iterator, List
@@ -18,6 +19,7 @@
 
 ROLE_IDS = ["managed-role-search-{}".format(uuid.uuid4()) for x in range(0, 500000)]
 USER_AUTH = {"username": "wikiuser", "password": "ujd_rbh5quw7GWC@pjc"}
+AUTH_BEARER = b64encode(("{}:{}".format(*USER_AUTH.values())).encode('ascii'))
 ROLE_TEMPLATE = {
     "indices": [
         {
@@ -143,7 +145,7 @@ def params(self):
             query = next(self._queries_iterator)
             return {
                 "method": "POST",
-                "headers": {"es-security-runas-user": USER_AUTH["username"]},
+                "headers": {"Authorization": "Basic {}".format(AUTH_BEARER)},
                 "path": f"{SEARCH_APPLICATION_ROOT_ENDPOINT}/{self.search_application_params.name}/_search",
                 "body": {
                     "params": {
@@ -183,6 +185,10 @@ async def create_users_and_roles(es, params):
     # For now we'll just work with one user with all the roles
     # num_users = params['users']
 
+    await es.security.put_user(
+        username=USER_AUTH["username"], params={"password": USER_AUTH["password"]}
+    )
+
     await es.indices.refresh(index="wikipedia")
     doc_count = await es.count(index="wikipedia")
 
@@ -208,17 +214,13 @@ async def create_users_and_roles(es, params):
         )
 
     await es.security.put_user(
-        username=USER_AUTH["username"], params={"roles": ROLE_IDS[0 : num_roles - 1], "password": USER_AUTH["password"]}
+        username=USER_AUTH["username"], params={"roles": ROLE_IDS[0 : num_roles - 1]}
     )
-    await es.security.put_role(name='impersonator', body={"run_as": ["wikiuser"]}, refresh="wait_for")
-    await es.security.put_user(username='esbench', params={"roles": ["superuser", "impersonator"]})
 
     await es.indices.refresh(index="wikipedia")
 
 
 async def reset_indices(es, params):
-    await es.security.delete_user(username=USER_AUTH["username"])
-
     roles = await es.security.get_role()
     for role in roles:
         name = list(role.keys())[0]
@@ -236,6 +238,10 @@ async def reset_indices(es, params):
 
     await es.indices.refresh(index="wikipedia")
 
+    await es.security.put_user(
+        username=USER_AUTH["username"], params={"roles": []}
+    )
+
 
 def register(registry):
     registry.register_param_source("query-string-search", QueryParamSource)