from burp import IBurpExtender
from burp import IContextMenuFactory
from burp import IBurpExtenderCallbacks
from burp import IHttpRequestResponse
from burp import IHttpListener

import zlib
import ast
from itsdangerous import base64_decode
from flask.sessions import SecureCookieSessionInterface


class MockApp(object):

    def __init__(self, secret_key):
        self.secret_key = secret_key


def session_cookie_encoder(secret_key, session_cookie_structure):
    """ Encode a Flask session cookie """
    try:
        app = MockApp(secret_key)

        session_cookie_structure = dict(ast.literal_eval(session_cookie_structure))
        si = SecureCookieSessionInterface()
        s = si.get_signing_serializer(app)

        return s.dumps(session_cookie_structure)
    except Exception as e:
        return "[Encoding error]{}".format(e)


def session_cookie_decoder(session_cookie_value, secret_key=None):
    """ Decode a Flask cookie  """
    try:
        if (secret_key == None):
            compressed = False
            payload = session_cookie_value

            if payload.startswith(b'.'):
                compressed = True
                payload = payload[1:]

            data = payload.split(".")[0]

            data = base64_decode(data)
            if compressed:
                data = zlib.decompress(data)

            return data
        else:
            app = MockApp(secret_key)

            si = SecureCookieSessionInterface()
            s = si.get_signing_serializer(app)

            return s.loads(session_cookie_value)
    except Exception as e:
        return "[Decoding error]{}".format(e)


def decode(cookie, secret_key):
    print(session_cookie_decoder(cookie, secret_key))


def encode(cookie_structure,secret_key):
    print(session_cookie_encoder(secret_key,cookie_structure))

from javax.swing import JMenuItem
from javax.swing import JOptionPane
class BurpExtender(IBurpExtender, IHttpListener, IContextMenuFactory, IHttpRequestResponse,
                   IBurpExtenderCallbacks):
    # 首先需要继承上面引入的接口,不然没法调用

    def registerExtenderCallbacks(self, callbacks):
        self._actionName = "decode_flask_session"  # 选项名称
        self._actionName1 = "encode_flask_session" # 选项名称
        self._helers = callbacks.getHelpers()
        self._callbacks = callbacks
        callbacks.setExtensionName("flask_session_manager") # 插件名称
        callbacks.registerHttpListener(self) # 注册
        callbacks.registerContextMenuFactory(self) # 注册

        return

    def createMenuItems(self, invocation):
    '''
    重写的方法，作用是创建菜单子选项，看函数名就差不多能理解了
    '''
        menu = []
        # invocation 是一个不知道哪来的值,没关系,文档里有这个,看看就懂了。
        # getSelectedMessages 也不知道哪来的,同上
        responses = invocation.getSelectedMessages()
        if len(responses) == 1:
            # 可以看到上面已经引用了java的swing包了,这就是jython的特性所在，可以引入java的包使用
            # 这里创建了两个子选项，按下动作分别指向两个函数 decode_session 和 encode_session
            menu.append(
                JMenuItem(self._actionName, None, actionPerformed=lambda x, inv=invocation: self.decode_session(inv)))
            menu.append(JMenuItem(self._actionName1,None,actionPerformed=lambda x: self.encode_session()))
            return menu
        return None

    def decode_session(self, invocation):
        # 解密函数，因为解密需要cookie和secret_key,所以这里造了一个对话框让用户输入已知的secret_key
        # 我不知道JOptionPane怎么用,于是百度了一下 'java swing JOptionPane' 就全有了
        secret_key = JOptionPane.showInputDialog(None, "secret_key:", "input", JOptionPane.QUESTION_MESSAGE)
        invMessage = invocation.getSelectedMessages()
        # request 就是一个完整的请求报文
        request = invMessage[0].getRequest().tostring()
        # 获取选中的需要解密的cookie所在的位置
        select_msg_index = invocation.getSelectionBounds()
        # 提取cookie
        select_cookie = request[select_msg_index[0]:select_msg_index[1]]
        # 调用上面的decode函数进行解密
        decode(select_cookie, secret_key)

    def encode_session(self):
        # cookie_structure是需要加密伪造的json格式的原始数据
        # 例如 {u'username': u'admin'}
        cookie_structure = JOptionPane.showInputDialog(None, "cookie_structure:", "input", JOptionPane.QUESTION_MESSAGE)
        secret_key = JOptionPane.showInputDialog(None, "secret_key:", "input", JOptionPane.QUESTION_MESSAGE)
        encode(cookie_structure, secret_key)