Include more headers with logins (#79)

* Include more headers with logins

* Fix lint

linkedin_matrix/commands/auth.py

@@ -1,4 +1,5 @@
 import logging
+import re
 
 from mautrix.bridge.commands import HelpSection, command_handler
 
@@ -83,7 +84,61 @@ async def login(evt: CommandEvent):
         return
 
     try:
-        await evt.sender.on_logged_in(cookies)
+        await evt.sender.on_logged_in(cookies, None)
+        await evt.reply("Successfully logged in")
+    except Exception as e:
+        logging.exception("Failed to log in")
+        await evt.reply(f"Failed to log in: {e}")
+        return
+
+
+@command_handler(
+    needs_auth=False,
+    management_only=False,
+    help_section=SECTION_AUTH,
+    help_text="""
+        Log in to LinkedIn using a "Copy as cURL" export from an existing LinkedIn browser session.
+    """,
+    help_args="<_curl command_>",
+)
+async def login_curl(evt: CommandEvent):
+    # if evt.sender.client and await evt.sender.client.logged_in():
+    #     await evt.reply("You're already logged in.")
+    #     return
+
+    if len(evt.args) == 0:
+        await evt.reply("**Usage:** `$cmdprefix+sp login_curl <cookie header>`")
+        return
+
+    # await evt.redact()
+
+    curl_command = " ".join(evt.args)
+
+    cookies: dict[str, str] = {}
+    headers: dict[str, str] = {}
+
+    curl_command_regex = r"-H '(?P<key>[^:]+): (?P<value>[^\']+)'"
+    header_matches = re.findall(curl_command_regex, curl_command)
+    for m in header_matches:
+        (name, value) = m
+
+        if name == "cookie":
+            cookie_items = value.split("; ")
+            for c in cookie_items:
+                n, v = c.split("=", 1)
+                cookies[n] = v
+        elif name == "accept":
+            # Every request will have a different value for this
+            pass
+        else:
+            headers[name] = value
+
+    if not cookies.get("li_at") or not cookies.get("JSESSIONID"):
+        await evt.reply("Missing li_at or JSESSIONID cookie")
+        return
+
+    try:
+        await evt.sender.on_logged_in(cookies, headers)
         await evt.reply("Successfully logged in")
     except Exception as e:
         logging.exception("Failed to log in")

linkedin_matrix/db/__init__.py

@@ -1,6 +1,7 @@
 from mautrix.util.async_db import Database
 
 from .cookie import Cookie
+from .http_header import HttpHeader
 from .message import Message
 from .model_base import Model
 from .portal import Portal
@@ -12,14 +13,15 @@
 
 
 def init(db: Database):
-    for table in (Cookie, Message, Portal, Puppet, Reaction, User, UserPortal):
+    for table in (HttpHeader, Cookie, Message, Portal, Puppet, Reaction, User, UserPortal):
         table.db = db  # type: ignore
 
 
 __all__ = (
     "init",
     "upgrade_table",
     # Models
+    "HttpHeader",
     "Cookie",
     "Message",
     "Model",

linkedin_matrix/db/http_header.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from asyncpg import Record
+from attr import dataclass
+
+from mautrix.types import UserID
+
+from .model_base import Model
+
+
+@dataclass
+class HttpHeader(Model):
+    mxid: UserID
+    name: str
+    value: str
+
+    _table_name = "http_header"
+    _field_list = [
+        "mxid",
+        "name",
+        "value",
+    ]
+
+    @classmethod
+    def _from_row(cls, row: Record | None) -> HttpHeader | None:
+        if row is None:
+            return None
+        return cls(**row)
+
+    @classmethod
+    async def get_for_mxid(cls, mxid: id.UserID) -> list[HttpHeader]:
+        query = HttpHeader.select_constructor("mxid=$1")
+        rows = await cls.db.fetch(query, mxid)
+        return [cls._from_row(row) for row in rows if row]
+
+    @classmethod
+    async def delete_all_for_mxid(cls, mxid: id.UserID):
+        await cls.db.execute("DELETE FROM http_header WHERE mxid=$1", mxid)
+
+    @classmethod
+    async def bulk_upsert(cls, mxid: id.UserID, http_headers: dict[str, str]):
+        for name, value in http_headers.items():
+            http_header = cls(mxid, name, value)
+            await http_header.upsert()
+
+    async def upsert(self):
+        query = """
+            INSERT INTO http_header (mxid, name, value)
+                 VALUES ($1, $2, $3)
+            ON CONFLICT (mxid, name)
+                DO UPDATE
+                    SET value=excluded.value
+        """
+        await self.db.execute(query, self.mxid, self.name, self.value)

linkedin_matrix/db/upgrade/__init__.py

@@ -12,6 +12,7 @@
     v07_puppet_contact_info_set,
     v08_splat_pickle_data,
     v09_cookie_table,
+    v10_http_header_table,
 )
 
 __all__ = (
@@ -24,4 +25,5 @@
     "v07_puppet_contact_info_set",
     "v08_splat_pickle_data",
     "v09_cookie_table",
+    "v10_http_header_table",
 )

linkedin_matrix/db/upgrade/v10_http_header_table.py

@@ -0,0 +1,18 @@
+from mautrix.util.async_db import Connection
+
+from . import upgrade_table
+
+
+@upgrade_table.register(description="Add a header table for storing all of the headers")
+async def upgrade_v10(conn: Connection):
+    await conn.execute(
+        """
+        CREATE TABLE http_header (
+            mxid  TEXT,
+            name  TEXT,
+            value TEXT,
+
+            PRIMARY KEY (mxid, name)
+        )
+        """
+    )

linkedin_matrix/user.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, AsyncGenerator, AsyncIterable, Awaitable, cast
+from typing import TYPE_CHECKING, AsyncGenerator, AsyncIterable, Awaitable, Optional, cast
 from asyncio.futures import Future
 from datetime import datetime
 import asyncio
@@ -26,7 +26,7 @@
 
 from . import portal as po, puppet as pu
 from .config import Config
-from .db import Cookie, User as DBUser
+from .db import Cookie, HttpHeader, User as DBUser
 
 if TYPE_CHECKING:
     from .__main__ import LinkedInBridge
@@ -195,7 +195,10 @@ async def load_session(self, is_startup: bool = False) -> bool:
             await self.push_bridge_state(BridgeStateEvent.BAD_CREDENTIALS, error="logged-out")
             return False
 
-        self.client = LinkedInMessaging.from_cookies({c.name: c.value for c in cookies})
+        self.client = LinkedInMessaging.from_cookies_and_headers(
+            {c.name: c.value for c in cookies},
+            {h.name: h.value for h in await HttpHeader.get_for_mxid(self.mxid)},
+        )
 
         backoff = 1.0
         while True:
@@ -255,10 +258,12 @@ async def is_logged_in(self) -> bool:
                 self.user_profile_cache = None
         return self._is_logged_in or False
 
-    async def on_logged_in(self, cookies: dict[str, str]):
+    async def on_logged_in(self, cookies: dict[str, str], headers: Optional[dict[str, str]]):
         cookies = {k: v.strip('"') for k, v in cookies.items()}
         await Cookie.bulk_upsert(self.mxid, cookies)
-        self.client = LinkedInMessaging.from_cookies(cookies)
+        if headers:
+            await HttpHeader.bulk_upsert(self.mxid, headers)
+        self.client = LinkedInMessaging.from_cookies_and_headers(cookies, headers)
         self.listener_event_handlers_created = False
         self.user_profile_cache = await self.client.get_user_profile()
         if (mp := self.user_profile_cache.mini_profile) and mp.entity_urn:

linkedin_matrix/web/provisioning_api.py

@@ -106,7 +106,7 @@ async def login(self, request: web.Request) -> web.Response:
             return web.HTTPBadRequest(body='{"error": "Missing keys"}', headers=self._headers)
 
         try:
-            await user.on_logged_in(data)
+            await user.on_logged_in(data, None)
             track(user, "$login_success")
         except Exception as e:
             track(user, "$login_failed", {"error": str(e)})