Handle Firefox containers in extract_token_from_browser.py (#909)

The extract_token_from_browser.py script does not handle Firefox containers
correctly. The container support works by namespacing cookies and local storage
with a 'userContextId'. In effect this means that when containers are used,
running the script will select the cookie from whichever container happens to be
returned first by the database query, which will be combined with the
non-containerised instance of the local storage.

Add correct handling by adding a new --container flag which specifies a
container either by numeric ID or by name, and including the correct
userContextId into the cookie query and local storage path (or excluding all
cookies from containers from the query if no --container flag is supplied).

With this the script can be run like:

python extract_token_from_browser.py --container Work firefox

to get the Slack auth token for the 'Work' container.

Signed-off-by: Toke Høiland-Jørgensen <toke@toke.dk>

1 files changed, 42 insertions, 7 deletions

diff --git a/extract_token_from_browser.py b/extract_token_from_browser.py
index 6e99e6b..d287f97 100755
--- a/extract_token_from_browser.py
+++ b/extract_token_from_browser.py
@@ -39,10 +39,12 @@ def sqlite3_connect(path: StrPath):
         con.close()
 
 
-def get_cookies(cookies_path: StrPath, cookie_query: str) -> tuple[str, str | None]:
+def get_cookies(
+    cookies_path: StrPath, cookie_query: str, params: tuple
+) -> tuple[str, str | None]:
     with sqlite3_connect(cookies_path) as con:
-        cookie_d_value = con.execute(cookie_query.format("d")).fetchone()
-        cookie_ds_value = con.execute(cookie_query.format("ds")).fetchone()
+        cookie_d_value = con.execute(cookie_query.format("d"), params).fetchone()
+        cookie_ds_value = con.execute(cookie_query.format("ds"), params).fetchone()
         if cookie_d_value and cookie_ds_value:
             return cookie_d_value[0], cookie_ds_value[0]
         elif cookie_d_value:
@@ -66,6 +68,12 @@ parser.add_argument(
 parser.add_argument(
     "--profile", help="Profile to look up cookies for", metavar="<profile>", nargs="?"
 )
+parser.add_argument(
+    "--container",
+    help="Firefox container to look up cookies for",
+    metavar="<id or name>",
+    nargs="?",
+)
 args = parser.parse_args()
 
 browser: Literal["firefox", "chrome"]
@@ -147,13 +155,40 @@ if browser == "firefox":
             sys.exit(1)
 
     cookies_path = default_profile_path.joinpath("cookies.sqlite")
+
+    if args.container:
+        try:
+            ctx_id = int(args.container)
+        except ValueError:
+            # non-numeric container ID, try to find by name
+            ctx_id = None
+            with open(default_profile_path.joinpath("containers.json"), "rb") as fp:
+                containers = json.load(fp)
+                for i in containers["identities"]:
+                    if "name" in i and i["name"] == args.container:
+                        ctx_id = i["userContextId"]
+                        break
+            if ctx_id is None:
+                print(
+                    f"Couldn't find Firefox container '{args.container}'",
+                    file=sys.stderr,
+                )
+                sys.exit(1)
+
+        userctx = f"^userContextId={ctx_id}"
+    else:
+        userctx = ""
+
     cookie_query = (
-        "SELECT value FROM moz_cookies WHERE host = '.slack.com' " "AND name = '{}'"
+        "SELECT value FROM moz_cookies WHERE originAttributes = ? "
+        "AND host = '.slack.com' AND name = '{}'"
+    )
+    cookie_d_value, cookie_ds_value = get_cookies(
+        cookies_path, cookie_query, (userctx,)
     )
-    cookie_d_value, cookie_ds_value = get_cookies(cookies_path, cookie_query)
 
     storage_path = default_profile_path.joinpath(
-        "storage/default/https+++app.slack.com/ls/data.sqlite"
+        f"storage/default/https+++app.slack.com{userctx}/ls/data.sqlite"
     )
     storage_query = "SELECT compression_type, conversion_type, value FROM data WHERE key = 'localConfig_v2'"
     local_config = None
@@ -195,7 +230,7 @@ elif browser == "chrome":
         "SELECT encrypted_value FROM cookies WHERE "
         "host_key = '.slack.com' AND name = '{}'"
     )
-    cookie_d_value, cookie_ds_value = get_cookies(cookies_path, cookie_query)
+    cookie_d_value, cookie_ds_value = get_cookies(cookies_path, cookie_query, ())
 
     bus = secretstorage.dbus_init()
     try: