snss.py

import enum

from construct import this, Bytes, Check, Computed, Enum, Flag, GreedyRange, Hex, If, Int16ul, Int32sl, Int32ul, Int64sl, Int8ul, Optional, PaddedString, Struct

class CommandId(enum.IntEnum):
    UpdateTabNavigation = 1
    RestoredEntry = 2
    WindowDeprecated = 3
    SelectedNavigationInTab = 4
    PinnedState = 5
    SetExtensionAppID = 6
    SetWindowAppName = 7
    SetTabUserAgentOverride = 8
    Window = 9
    SetTabGroupData = 10
    SetTabUserAgentOverride2 = 11
    SetWindowUserTitle = 12
    CreateGroup = 13
    AddTabExtraData = 14
    InitialStateMarker = 255

SessionId = Int32sl

String = Struct(
    'length' / Int32sl,
    Check(this.length > 0),
    'data' / PaddedString(this.length, 'utf-8')
)
String16 = Struct(
    'length' / Int32sl,
    Check(this.length > 0),
    'data' / PaddedString(this.length*2, 'utf-16')
)

array_to_map = lambda obj,_: {item.key: item.value for item in obj if item.key is not None and item.value is not None}

ExtendedInfoMap = Struct(
    'size' / Int32sl,
    'items' / If(this.size > 0,
        Struct(
            'key' / Optional(String),
            'value' / Optional(String),
        )[this.size] * array_to_map
    )
)

def Pickle(*args):
    return Struct(
        'size' / Int32ul,
        '_raw' / Bytes(this.size),
        'contents' / Computed(lambda this: Struct(*args).parse(this._raw)),
    )

UpdateTabNavigationPayload = Pickle(
    'id' / SessionId,
    'index' / Int32sl,
    'virtual_url_spec' / String,
    'title' / String16,
    'encoded_page_state' / String,
    'transition_type' / Int32sl,
    'type_mask' / Optional(Int32ul),
    'opt' / If(this.type_mask is not None,
        Struct(
            'referrer_spec' / Optional(String),
            '_referrer_policy' / Optional(Int32sl),  # obsolete
            'original_request_url_spec' / Optional(String),
            'is_overriding_user_agent' / Optional(Flag),
            'timestamp' / Optional(Int64sl),
            '_search_terms' / Optional(String16),  # obsolete
            'http_status_code' / Optional(Int32sl),
            'referrer_policy' / Optional(Int32sl),
            'extended_info_map' / Optional(ExtendedInfoMap),
            'task_id' / Optional(Int64sl),
            'parent_task_id' / Optional(Int64sl),
            'root_task_id' / Optional(Int64sl),
            '_children_task_ids_size' / Optional(Int32sl),
        )
    )
)

SelectedNavigationInTabPayload = Struct(
    'id' / SessionId,
    'index' / Int32sl,
    'timestamp' / Optional(Int64sl)
)

def command_payload(cmd):
    payload_type = {
        'UpdateTabNavigation': UpdateTabNavigationPayload,
        'SelectedNavigationInTab': SelectedNavigationInTabPayload,
    }.get(cmd.id)
    if payload_type is None:
        return None
    try:
        return payload_type.parse(cmd._raw)
    except Exception as e:
        print(e)
        raise

Command = Struct(
    'size' / Int16ul,
    'id' / Enum(Int8ul, CommandId),
    '_raw' / Bytes(this.size-1),
    'payload' / Computed(lambda this: command_payload(this)),
)

Header = Struct(
    'signature' / Hex(Int32ul),
    'version' / Int32ul
)

SessionFile = Struct(
    'header' / Header,
    'commands' / GreedyRange(Command)
)

def parse_session_data(data):
    session = SessionFile.parse(data)
    parsed = Header.sizeof() + sum(2+cmd.size for cmd in session.commands)
    if parsed != len(data):
        print(f'Failed to fully parse session data! Parsed {parsed}/{len(data)} bytes.')
    return session