table for cookies

Signed-off-by: Sumner Evans <[email protected]>

linkedin_matrix/commands/auth.py

@@ -1,3 +1,4 @@
+import json
 import logging
 
 from mautrix.bridge.commands import HelpSection, command_handler
@@ -54,26 +55,36 @@ async def whoami(evt: CommandEvent):
     needs_auth=False,
     management_only=False,
     help_section=SECTION_AUTH,
-    help_text=(
-        "Log in to LinkedIn by cookies from an existing LinkedIn browser session (recommended "
-        "to use a private window to extract the cookies)"
-    ),
-    help_args="<_li\\_at_> <_jsessionid_>",
+    help_text="""
+        Log in to LinkedIn using cookies from an existing LinkedIn browser session. To extract the
+        cookies go to your browser developer tools, open the Network tab, then copy the `Cookie`
+        header from one of the requests to `https://www.linkedin.com/` and paste the result into
+        the command. It is recommended that you use a private window to extract the cookies.
+    """,
+    help_args="<_cookie header_>",
 )
 async def login(evt: CommandEvent):
     if evt.sender.client and await evt.sender.client.logged_in():
         await evt.reply("You're already logged in.")
         return
-    elif len(evt.args) != 2:
-        await evt.reply("**Usage:** `$cmdprefix+sp login <li_at> <jsessionid>`")
+
+    if len(evt.args) == 0:
+        await evt.reply("**Usage:** `$cmdprefix+sp login <cookie header>`")
         return
 
-    li_at = evt.args[0].strip('"')
-    jsessionid = evt.args[1].strip('"')
     await evt.redact()
 
+    cookies: dict[str, str] = {}
+    for cookie in evt.args:
+        key, val = cookie.strip(" ;").split("=", 1)
+        cookies[key] = val
+
+    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(li_at, jsessionid)
+        await evt.sender.on_logged_in(cookies)
         await evt.reply("Successfully logged in")
     except Exception as e:
         logging.exception("Failed to log in")

linkedin_matrix/db/__init__.py

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

linkedin_matrix/db/cookie.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from typing import cast
+
+from asyncpg import Record
+from attr import dataclass
+
+from linkedin_messaging import URN
+from mautrix.types import RoomID, UserID
+
+from .model_base import Model
+
+
+@dataclass
+class Cookie(Model):
+    mxid: UserID
+    name: string
+    value: string
+
+    _table_name = "cookie"
+    _field_list = [
+        "mxid",
+        "name",
+        "value",
+    ]
+
+    @classmethod
+    def _from_row(cls, row: Record | None) -> Cookie | None:
+        if row is None:
+            return None
+        return cls(**row)
+
+    @classmethod
+    async def get_for_mxid(cls, mxid: id.UserID) -> list[Cookie]:
+        query = Cookie.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 self.db.execute("DELETE FROM cookie WHERE mxid=$1", self.mxid)
+
+    @classmethod
+    async def bulk_upsert(cls, mxid: id.UserID, cookies: dict[str, str]):
+        for name, value in cookies.items():
+            cookie = cls(mxid, name, value)
+            await cookie.upsert()
+
+    async def upsert(self):
+        query = """
+            INSERT INTO cookie (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

@@ -11,6 +11,7 @@
     v06_add_space_mxid_to_user,
     v07_puppet_contact_info_set,
     v08_splat_pickle_data,
+    v09_cookie_table,
 )
 
 __all__ = (
@@ -22,4 +23,5 @@
     "v06_add_space_mxid_to_user",
     "v07_puppet_contact_info_set",
     "v08_splat_pickle_data",
+    "v09_cookie_table",
 )

linkedin_matrix/db/upgrade/v09_cookie_table.py

@@ -0,0 +1,39 @@
+import logging
+import pickle
+
+from mautrix.util.async_db import Connection
+
+from . import upgrade_table
+
+
+@upgrade_table.register(description="Add a cookie table for storing all of the cookies")
+async def upgrade_v9(conn: Connection):
+    await conn.execute(
+        """
+        CREATE TABLE cookie (
+            mxid  TEXT,
+            name  TEXT,
+            value TEXT,
+
+            PRIMARY KEY (mxid, name)
+        )
+        """
+    )
+
+    for row in await conn.fetch('SELECT mxid, jsessionid, li_at FROM "user"'):
+        mxid = row["mxid"]
+        jsessionid = row["jsessionid"]
+        li_at = row["li_at"]
+
+        if jsessionid:
+            await conn.execute(
+                "INSERT INTO cookie (mxid, name, value) VALUES ($1, 'JSESSIONID', $2)",
+                mxid,
+                jsessionid,
+            )
+        if li_at:
+            await conn.execute(
+                "INSERT INTO cookie (mxid, name, value) VALUES ($1, 'li_at', $2)",
+                mxid,
+                li_at,
+            )

linkedin_matrix/db/user.py

@@ -18,15 +18,10 @@ class User(Model):
     notice_room: RoomID | None
     space_mxid: RoomID | None
 
-    jsessionid: str | None
-    li_at: str | None
-
     _table_name = "user"
     _field_list = [
         "mxid",
         "li_member_urn",
-        "jsessionid",
-        "li_at",
         "notice_room",
         "space_mxid",
     ]
@@ -66,8 +61,6 @@ async def insert(self):
             query,
             self.mxid,
             self.li_member_urn.id_str() if self.li_member_urn else None,
-            self.jsessionid,
-            self.li_at,
             self.notice_room,
             self.space_mxid,
         )
@@ -79,18 +72,14 @@ async def save(self):
         query = """
             UPDATE "user"
                SET li_member_urn=$2,
-                   jsessionid=$3,
-                   li_at=$4,
-                   notice_room=$5,
-                   space_mxid=$6
+                   notice_room=$3,
+                   space_mxid=$4
              WHERE mxid=$1
         """
         await self.db.execute(
             query,
             self.mxid,
             self.li_member_urn.id_str() if self.li_member_urn else None,
-            self.jsessionid,
-            self.li_at,
             self.notice_room,
             self.space_mxid,
         )

linkedin_matrix/user.py

@@ -26,7 +26,7 @@
 
 from . import portal as po, puppet as pu
 from .config import Config
-from .db import User as DBUser
+from .db import Cookie, User as DBUser
 
 if TYPE_CHECKING:
     from .__main__ import LinkedInBridge
@@ -62,15 +62,11 @@ def __init__(
         self,
         mxid: UserID,
         li_member_urn: URN | None = None,
-        jsessionid: str | None = None,
-        li_at: str | None = None,
         notice_room: RoomID | None = None,
         space_mxid: RoomID | None = None,
     ):
-        super().__init__(mxid, li_member_urn, notice_room, space_mxid, jsessionid, li_at)
+        super().__init__(mxid, li_member_urn, notice_room, space_mxid)
         BaseUser.__init__(self)
-        self.notice_room = notice_room
-        self.space_mxid = space_mxid
         self._notice_room_lock = asyncio.Lock()
         self._notice_send_lock = asyncio.Lock()
 
@@ -193,18 +189,21 @@ async def get_portal_with(self, puppet: pu.Puppet, create: bool = True) -> po.Po
     async def load_session(self, is_startup: bool = False) -> bool:
         if self._is_logged_in and is_startup:
             return True
-        if not self.jsessionid or not self.li_at:
+        cookies = await Cookie.get_for_mxid(self.mxid)
+        cookie_names = set(c.name for c in cookies)
+        if "li_at" not in cookie_names or "JSESSIONID" not in cookie_names:
             await self.push_bridge_state(BridgeStateEvent.BAD_CREDENTIALS, error="logged-out")
             return False
 
-        self.client = LinkedInMessaging.from_cookies(self.li_at, self.jsessionid)
+        self.client = LinkedInMessaging.from_cookies({c.name: c.value for c in cookies})
 
         backoff = 1.0
         while True:
             try:
                 self.user_profile_cache = await self.client.get_user_profile()
                 break
             except (TooManyRedirects, ServerConnectionError) as e:
+                self.log.info(f"Failed to get user profile: {e}")
                 await self.push_bridge_state(BridgeStateEvent.BAD_CREDENTIALS, message=str(e))
                 return False
             except Exception as e:
@@ -256,10 +255,10 @@ 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, li_at: str, jsessionid: str):
-        self.li_at = li_at
-        self.jsessionid = jsessionid
-        self.client = LinkedInMessaging.from_cookies(self.li_at, self.jsessionid)
+    async def on_logged_in(self, cookies: 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)
         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:
@@ -307,10 +306,9 @@ async def logout(self):
                 del self.by_li_member_urn[self.li_member_urn]
             except KeyError:
                 pass
+        await Cookie.delete_all_for_mxid(self.mxid)
         self._track_metric(METRIC_LOGGED_IN, True)
         self.client = None
-        self.jsessionid = None
-        self.li_at = None
         self.listener_event_handlers_created = False
         self.user_profile_cache = None
         self.li_member_urn = None

linkedin_matrix/web/provisioning_api.py

@@ -80,7 +80,7 @@ async def status(self, request: web.Request) -> web.Response:
             user_profile = user.user_profile_cache
             if user_profile is not None:
                 self.log.debug("Cache hit on user_profile_cache")
-            user_profile = user_profile or await user.client.get_user_profile()
+            user_profile = user_profile or await user.client.get_user_metadata()
             data["linkedin"] = user_profile.to_dict()
 
         return web.json_response(data, headers=self._acao_headers)
@@ -101,8 +101,7 @@ async def login(self, request: web.Request) -> web.Response:
             return web.HTTPBadRequest(body='{"error": "Missing keys"}', headers=self._headers)
 
         try:
-            jsessionid = data["JSESSIONID"].strip('"')
-            await user.on_logged_in(data["li_at"], jsessionid)
+            await user.on_logged_in(data)
             track(user, "$login_success")
         except Exception as e:
             track(user, "$login_failed", {"error": str(e)})

linkedin_messaging/linkedin.py

@@ -139,12 +139,15 @@ def __init__(self):
         self.event_listeners = defaultdict(list)
 
     @staticmethod
-    def from_cookies(li_at: str, jsessionid: str) -> "LinkedInMessaging":
+    def from_cookies(cookies: dict[str, str]) -> "LinkedInMessaging":
         linkedin = LinkedInMessaging()
-        linkedin.session.cookie_jar.update_cookies({"li_at": li_at, "JSESSIONID": jsessionid})
-        linkedin._request_headers["csrf-token"] = jsessionid
+        linkedin.session.cookie_jar.update_cookies(cookies)
+        linkedin._request_headers["csrf-token"] = cookies["JSESSIONID"].strip('"')
         return linkedin
 
+    def cookies(self) -> dict[str, str]:
+        return {c.key: c.value for c in self.session.cookie_jar}
+
     async def close(self):
         await self.session.close()
 
@@ -174,13 +177,13 @@ async def logged_in(self) -> bool:
             logging.exception(f"Failed getting the user profile: {e}")
             return False
 
-    async def login_manual(self, li_at: str, jsessionid: str, new_session: bool = True):
+    async def login_manual(self, cookies: dict[str, str], new_session: bool = True):
         if new_session:
             if self.session:
                 await self.session.close()
             self.session = aiohttp.ClientSession()
-        self.session.cookie_jar.update_cookies({"li_at": li_at, "JSESSIONID": jsessionid})
-        self._request_headers["csrf-token"] = jsessionid.strip('"')
+        self.session.cookie_jar.update_cookies(cookies)
+        self._request_headers["csrf-token"] = cookies["JSESSIONID"].strip('"')
 
     async def login(self, email: str, password: str, new_session: bool = True):
         if new_session:
@@ -417,7 +420,6 @@ async def send_message(
                 f"/messaging/conversations/{conversation_id}/events",
                 params=params,
                 json=message_event,
-                headers=self._request_headers,
             )
 
         return cast(SendMessageResponse, await try_from_json(SendMessageResponse, res))