'''
    SNSS Session File Parser for PyBrinf
    This module parses SNSS files and returns a list of current tabs.
    Script made by Manuel Cabral
'''

import os, struct
from io import StringIO, BytesIO
from enum import Enum, EnumMeta

class MetaEnum(EnumMeta):
    '''This class is used for creating an enum with a custom __contains__ method'''
    def __contains__(cls, item: int) -> bool:
        try:
            cls(item)
        except ValueError:
            return False
        return True    

class SNSSTypeCommand(Enum, metaclass=MetaEnum):
    '''Enum for SNSS command types'''
    CommandSetTabWindow = 0
    CommandSetTabIndexInWindow = 2
    CommandTabClosed = 3
    CommandWindowClosed = 4
    CommandTabNavigationPathPrunedFromBack = 5
    CommandUpdateTabNavigation = 6
    CommandSetSelectedNavigationIndex = 7
    CommandSetSelectedTabInIndex = 8
    CommandSetWindowType = 9
    CommandTabNavigationPathPrunedFromFront = 11
    CommandSetPinnedState = 12
    CommandSetExtensionAppID = 13
    CommandSetWindowBounds3 = 14


class SNSSCommand:
    '''Command class core'''
    def __init__(self, id: int, content: bytes):
        self.id = id
        self.content = content

def parse_commands(commands):
    output = []
    for command in commands:
        if command.id in SNSSTypeCommand:
            if command.id == SNSSTypeCommand.CommandUpdateTabNavigation.value:

                # Extracting the pickle payload
                content = BytesIO(command.content)

                # Get the size of the pickle
                content.seek(0, os.SEEK_END)
                pickle_size = content.tell()
                content.seek(0, os.SEEK_SET)

                # Get the size of the payload, unpacking uint32
                payload_size, = struct.unpack('I', content.read(4))
                payloard_start = pickle_size - payload_size

                # Get the payload, unpacking uint32
                tab_id, = struct.unpack('I', content.read(4))
                index, = struct.unpack('I', content.read(4))

                # Get the url size, unpacking uint32
                str_length, = struct.unpack('I', content.read(4))
                # If the url size is more than the difference between 
                # the pickle size and the current position, it's invalid
                if str_length > pickle_size - content.tell():
                    raise Exception("Invalid string length")
                url = content.read(str_length).decode('utf-8', 'ignore')

                output.append((tab_id, index, url))
    return output

def parse(path):
    # Parse a SNSS Session File to get all commands
    commands = []
    file = open(path, 'rb')

    # Getting file size
    file.seek(0, os.SEEK_END)
    end = file.tell()
    file.seek(0, os.SEEK_SET)

    # Reading signature, unpacking int32
    signature, = struct.unpack('i', file.read(4))
    if signature != 0x53534E53: # SNSS
        raise Exception("Invalid SNSS file")
        
    # Reading version version, unpacking int32
    version, = struct.unpack('i', file.read(4))

    while end - file.tell() > 0:
        # Read command size, unpacking uint16
        command_size, = struct.unpack('H', file.read(2))
        if command_size == 0:
            raise Exception("Invalid command size, maybe corrupted file")
        
        # Read command id, unpacking uint8
        id, = struct.unpack('B', file.read(1))

        # Read content, NOTE: command id is included in command_size
        content = file.read(command_size - 1)

        command = SNSSCommand(id, content)
        commands.append(command)

    file.close()
    return commands

if __name__ == '__main__':
    commands = parse('Session')
    # for now only gets the tab navigation
    tabs = parse_commands(commands)
    for tab in tabs:
        print(tab)