#!/usr/bin/python3 -sP
# GUI for smbcmp
#
# Copyright (C) 2019 Mairo Paul Rufus <akoudanilo@gmail.com>
# Copyright (C) 2019 Aurelien Aptel <aurelien.aptel@gmail.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


import re
import functools

import wx
import wx.stc
import wx.adv
import wx.lib.mixins.listctrl
import smbcmp

DIFFLIB = 1
PDML = 2

ARGS = None

def main():
    global ARGS
    ARGS = smbcmp.parse_args(gui=True)
    app = Smbcmp(0)
    app.MainLoop()

class TracePanel(wx.Panel):
    """ Panel holding the summaries of the packets """
    def __init__(self, *args, **kwds):
        kwds["style"] = kwds.get("style", 0)
        wx.Panel.__init__(self, *args, **kwds)
        self.cap = ""
        self.nos = [1]
        self.data = []
        self.buffer = wx.ListBox(self, wx.ID_ANY)
        self._reset_items()
        self.debug = []
        self.__set_properties()
        self.__do_layout()

    def __set_properties(self):
        pass

    def __do_layout(self):
        sizer_1 = wx.BoxSizer(wx.HORIZONTAL)
        sizer_1.Add(self.buffer, 1, wx.ALL | wx.EXPAND, 5)
        self.SetSizer(sizer_1)
        sizer_1.Fit(self)
        self.Layout()

    def _reset_items(self):
        self.buffer.Clear()

        if self.cap == "":
            self.buffer.InsertItems(items=["No capture opened"], pos=0)
            return

        if len(self.data) == 0:
            self.buffer.InsertItems(items=["No SMB packets"], pos=0)
            return

        self.buffer.InsertItems(items=self.data, pos=0)
        self.buffer.SetSelection(0)

    def set_data(self, file_path):
        """ Fill the Panel with smb summaries """
        self.cap = file_path
        pkts = smbcmp.smb_summaries(self.cap)
        self.nos = sorted(pkts.keys())
        self.data = [pkts[x] for x in self.nos]
        self._reset_items()

    @property
    def cap(self):
        return self._cap

    @cap.setter
    def cap(self, value):
        self._cap = value

    @property
    def data(self):
        return self._data

    @data.setter
    def data(self, value):
        self._data = value

    def get_selected_packet(self):
        sel = self.buffer.GetSelection()
        if self.cap and len(self.data) > 0 and sel >= 0:
            return (self.cap, self.nos[sel])
        return None

    def permute_data(self, panel):
        sel_self = self.buffer.GetSelection()
        sel_panel = panel.buffer.GetSelection()

        # swap user-data
        self.nos, panel.nos = panel.nos, self.nos
        self.data, panel.data = panel.data, self.data
        self.cap, panel.cap = panel.cap, self.cap

        # reset items and selection
        self._reset_items()
        panel._reset_items()
        if sel_panel >= 0:
            self.buffer.SetSelection(sel_panel)
        if sel_self >= 0:
            panel.buffer.SetSelection(sel_self)


    def move(self, direction):
        selected = self.buffer.GetSelection()
        if direction + 1:   # From top to bottom
            self.buffer.SetSelection((selected + 1) % len(self.data))
        elif selected > 0:
            self.buffer.SetSelection(selected - 1)


class MainFrame(wx.Frame):
    def __init__(self, *args, **kwds):
        kwds["style"] = kwds.get("style", 0) | wx.DEFAULT_FRAME_STYLE
        wx.Frame.__init__(self, *args, **kwds)

        self.lpanel = TracePanel(self, wx.ID_ANY)
        self.rpanel = TracePanel(self, wx.ID_ANY)

        self.data = ["diffs"]
        self.bbuffer = DiffCtrl(self, wx.ID_ANY)
        self.__set_properties()
        self.__do_layout()
        self.SetSize(1000,800)
        self.create_menu()

        self.rpanel.buffer.Bind(wx.EVT_KEY_DOWN, self.on_key_press)
        self.rpanel.buffer.Bind(wx.EVT_LISTBOX, self.on_summary_change)
        self.lpanel.buffer.Bind(wx.EVT_KEY_DOWN, self.on_key_press)
        self.lpanel.buffer.Bind(wx.EVT_LISTBOX, self.on_summary_change)
        self.bbuffer.Bind(wx.EVT_KEY_DOWN, self.on_key_press)

        if ARGS.filea:
            self.lpanel.set_data(ARGS.filea)
        if ARGS.fileb:
            self.rpanel.set_data(ARGS.fileb)

        if ARGS.mode == 'text':
            self.bbuffer.diff_method = DIFFLIB
        elif ARGS.mode == 'pdml':
            self.bbuffer.diff_method = PDML

        if ARGS.filea and ARGS.fileb:
            self.refresh()

    def on_summary_change(self, event):
        self.refresh()

    def on_key_press(self, event):
        keycode = event.GetKeyCode()

        if keycode == wx.WXK_DOWN:
            self.rpanel.move(+1)
            self.lpanel.move(+1)
        elif keycode == wx.WXK_UP:
            self.rpanel.move(-1)
            self.lpanel.move(-1)

        elif keycode == ord("D"):
            self.lpanel.move(+1)
        elif keycode == ord("F"):
            self.lpanel.move(-1)

        elif keycode == ord("J"):
            self.rpanel.move(+1)
        elif keycode == ord("K"):
            self.rpanel.move(-1)

        elif keycode == ord("B"):
            self.bbuffer.move(+1)
        elif keycode == ord("N"):
            self.bbuffer.move(-1)

        else:
            event.Skip()
            return

        self.refresh()

    def refresh(self):
        if not self.lpanel.cap or not self.rpanel.cap:
            return

        pkt_a = self.lpanel.get_selected_packet()
        pkt_b = self.rpanel.get_selected_packet()

        if not pkt_a or not pkt_b:
            return

        if self.bbuffer.diff_method == DIFFLIB:
            self.bbuffer.refresh_difflib(smbcmp.smb_diff(pkt_a, pkt_b, gui=True))
        elif self.bbuffer.diff_method == PDML:
            self.bbuffer.refresh_pdml(smbcmp.PDMLDiff(pkt_a, pkt_b).smb_diff())
        else:
            raise Exception("Unknown method")

    def create_menu(self):
        menu_bar = wx.MenuBar()

        file_menu = wx.Menu()
        open_file1_menu_item = file_menu.Append(
            wx.ID_ANY, 'Open Capture &1',
            'Open the first capture file'
        )
        open_file2_menu_item = file_menu.Append(
            wx.ID_ANY, 'Open Capture &2',
            'Open the second capture file'
        )
        permute_menu_item = file_menu.Append(
            wx.ID_ANY, '&Permute the two captures',
            'Permute the two captures')
        close_menu_item = file_menu.Append(
            wx.ID_ANY, "&Close",
            "Close the application")
        menu_bar.Append(file_menu, '&File')
        self.Bind(
            event=wx.EVT_MENU,
            handler=self.on_open_file1,
            source=open_file1_menu_item,
        )
        self.Bind(
            event=wx.EVT_MENU,
            handler=self.on_open_file2,
            source=open_file2_menu_item)
        self.Bind(
            event=wx.EVT_MENU,
            handler=self.on_permute,
            source=permute_menu_item)
        self.Bind(
            event=wx.EVT_MENU,
            handler=self.on_close,
            source=close_menu_item
        )

        options_menu = wx.Menu()
        engine_menu_item = options_menu.Append(
            wx.ID_ANY, "&Engine",
            "Configure the application")
        colors_menu_item = options_menu.Append(
            wx.ID_ANY, "&Colors",
            "Change the color scheme")
        menu_bar.Append(options_menu, '&Options')
        self.Bind(
            event=wx.EVT_MENU,
            handler=self.on_engine,
            source=engine_menu_item)
        self.Bind(
            event=wx.EVT_MENU,
            handler=self.on_colors,
            source=colors_menu_item)

        help_menu = wx.Menu()
        about_menu_item = help_menu.Append(
            wx.ID_ANY, "&About",
            "Open the about Box")
        menu_bar.Append(help_menu, '&Help')
        self.Bind(
            event=wx.EVT_MENU,
            handler=self.on_about,
            source=about_menu_item)

        self.SetMenuBar(menu_bar)

    def on_colors(self, event):
        dlg1 = wx.ColourDialog(self)
        if dlg1.ShowModal() == wx.ID_OK:
            color = dlg1.GetColourData().GetColour()
            self.bbuffer.color_add = color
        dlg2 = wx.ColourDialog(self)
        if dlg2.ShowModal() == wx.ID_OK:
            color = dlg2.GetColourData().GetColour()
            self.bbuffer.color_rem = color

    def on_engine(self, event):
        dlg = ChooseEngineDialog(None,
                                 title='Choose a diffs engine')
        dlg.ShowModal()
        self.bbuffer.diff_method = dlg.choice
        dlg.Destroy()

    def on_open_file1(self, event):
        title = "Choose a capture file:"
        dlg = wx.FileDialog(self,
                            message=title,
                            defaultDir="~",
                            style=wx.DD_DEFAULT_STYLE)
        if dlg.ShowModal() == wx.ID_OK:
            self.lpanel.set_data(dlg.GetPath())
        dlg.Destroy()

    def on_open_file2(self, event):
        title = "Choose a capture file:"
        dlg = wx.FileDialog(self,
                            message=title,
                            defaultDir="~",
                            style=wx.DD_DEFAULT_STYLE)
        if dlg.ShowModal() == wx.ID_OK:
            self.rpanel.set_data(dlg.GetPath())
        dlg.Destroy()

    def on_permute(self, event):
        self.lpanel.permute_data(self.rpanel)
        self.refresh()

    def on_close(self, event):
        self.Close()

    def on_about(self, event):
        info = wx.adv.AboutDialogInfo()
        info.Name = "About"
        info.Version = "smbcmp-gui v0.1"
        info.Copyright = ("(C) 2019 Mairo Paul Rufus <akoudanilo@gmail.com>\n"
                          +"Aurelien Aptel <aurelien.aptel@gmail.com>")
        info.Description = "Diff, compare and debug SMB traffic. This is the GUI version of smbcmp."
        info.WebSite = ("https://smbcmp.github.io", "Project homepage")
        info.Developers = ["RMPR", "Aurelien Aptel"]
        info.License = "GNU General Public License V3"
        wx.adv.AboutBox(info)

    def __set_properties(self):
        self.SetTitle("smbcmp")

    def __do_layout(self):
        main_sizer = wx.BoxSizer(wx.VERTICAL)
        sizer_top = wx.BoxSizer(wx.HORIZONTAL)
        sizer_top.Add(self.lpanel, 1, wx.EXPAND, 0)
        sizer_top.Add(self.rpanel, 1, wx.EXPAND, 0)
        main_sizer.Add(sizer_top, 1, wx.EXPAND, 0)
        main_sizer.Add(self.bbuffer, 1, wx.ALL | wx.EXPAND, 15)
        self.SetSizer(main_sizer)
        main_sizer.Fit(self)
        self.Layout()

class DiffCtrl(wx.Panel):
    def __init__(self, *args, **kwargs):
        super(DiffCtrl, self).__init__(*args, **kwargs)
        self.diff_method = DIFFLIB
        self.txt = wx.stc.StyledTextCtrl(self, wx.ID_ANY, style=wx.EXPAND|wx.TE_MULTILINE|wx.TE_READONLY|wx.TE_DONTWRAP|wx.TE_RICH)

        #
        # styling
        #
        self.style_normal = 0 # default
        self.style_add = 2
        self.style_rem = 3
        self.color_add = wx.Colour(0, 168, 107)  # red
        self.color_rem = wx.Colour(255, 127, 127)  # green
        self.color_mod = wx.Colour(0, 135, 189) # blue
        self.font = wx.Font(wx.FontInfo().Family(wx.FONTFAMILY_TELETYPE))

        #self.txt.StyleSetBackground(self.style_normal, wx.BLACK)
        #self.txt.StyleSetEOLFilled(self.style_normal, True)
        self.txt.StyleSetFont(self.style_normal, self.font)
        self.txt.StyleSetBackground(self.style_add, self.color_add)
        self.txt.StyleSetEOLFilled(self.style_add, True)
        self.txt.StyleSetFont(self.style_add, self.font)
        self.txt.StyleSetBackground(self.style_rem, self.color_rem)
        self.txt.StyleSetEOLFilled(self.style_rem, True)
        self.txt.StyleSetFont(self.style_rem, self.font)

        self.txt.SetReadOnly(True)
        self.txt.SetEditable(False)

        bsizer = wx.BoxSizer()
        bsizer.Add(self.txt, 1, wx.EXPAND)
        self.SetSizerAndFit(bsizer)

    @property
    def diff_method(self):
        return self._diff_method

    @diff_method.setter
    def diff_method(self, value):
        self._diff_method = value

    @property
    def color_add(self):
        return self._color_add

    @color_add.setter
    def color_add(self, value):
        self.txt.StyleSetBackground(self.style_add, value)
        self._color_add = value

    @property
    def color_rem(self):
        return self._color_rem

    @color_rem.setter
    def color_rem(self, value):
        self.txt.StyleSetBackground(self.style_rem, value)
        self._color_rem = value

    def move(self, *args):
        pass

    def _append_text(self, text, style=None):
        start = self.txt.GetLength()
        self.txt.AppendText(text)
        nb = self.txt.GetLength() - start
        if style:
            self.txt.StartStyling(start, 31)
            self.txt.SetStyling(nb, style)

    def refresh_difflib(self, data):
        self.txt.SetEditable(True)
        self.txt.SetReadOnly(False)
        self.txt.ClearAll()

        for line in data:
            style = None
            if line[0] == "+":
                style = self.style_add
            elif line[0] == "-":
                style = self.style_rem

            self._append_text(line+"\n", style)

        self.txt.SetEditable(False)
        self.txt.SetReadOnly(True)

    def refresh_pdml(self, data):
        self.txt.SetEditable(True)
        self.txt.SetReadOnly(False)
        self.txt.ClearAll()

        t_newline = 1
        t_text = 2
        t_esc = 3
        rx = r'(\n)|([^\x1b\n]+)|(\x1b\[.+?m)'

        # TODO: use DiffOutputItem data for better formating (tree-like structure)
        data = data.get_text()

        style = None
        for tok in re.finditer(rx, data):
            typ = tok.lastindex
            val = tok.group()
            if typ == t_text:
                self._append_text(val, style)
            elif typ == t_newline:
                self._append_text("\n", style)
            elif typ == t_esc:
                assert(val[-1] == 'm')
                for num in val[2:-1].split(';'):
                    num = int(num)
                    if num == 0:
                        style = None
                    elif num == 30:
                        style = None
                    elif num == 31:
                        style = self.style_rem
                    elif num == 32:
                        style = self.style_add
                    elif num == 1:
                        # ignore for now
                        pass
                    else:
                        raise Exception("unsupported CSI color")
            else:
                raise Exception("invalid token %d" % typ)

        self.txt.SetEditable(False)
        self.txt.SetReadOnly(True)


class ChooseEngineDialog(wx.Dialog):
    """ Dialog box which allows to choose diffing methods"""

    def __init__(self, *args, **kwargs):
        super(ChooseEngineDialog, self).__init__(*args, **kwargs)

        self.choice = DIFFLIB
        self.InitUI()
        self.SetSize((250, 200))
        self.SetTitle("Choose Diffs Engine")

    def InitUI(self):
        pnl = wx.Panel(self)
        vbox = wx.BoxSizer(wx.VERTICAL)

        sb = wx.StaticBox(pnl, label='Engines')
        sbs = wx.StaticBoxSizer(sb, orient=wx.VERTICAL)

        self.radio = wx.RadioButton(pnl, label='Difflib', style=wx.RB_GROUP)
        self.radio2 = wx.RadioButton(pnl, label='Pdml')
        sbs.Add(self.radio)
        sbs.Add(self.radio2)

        pnl.SetSizer(sbs)

        hbox2 = wx.BoxSizer(wx.HORIZONTAL)
        okButton = wx.Button(self, label='Ok')
        closeButton = wx.Button(self, label='Close')
        hbox2.Add(okButton)
        hbox2.Add(closeButton, flag=wx.LEFT, border=5)

        vbox.Add(pnl, proportion=1, flag=wx.ALL | wx.EXPAND, border=5)
        vbox.Add(hbox2, flag=wx.ALIGN_CENTER | wx.TOP | wx.BOTTOM, border=10)

        self.SetSizer(vbox)

        okButton.Bind(wx.EVT_BUTTON, self.on_btn_ok)
        closeButton.Bind(wx.EVT_BUTTON, self.on_close)

    def on_btn_ok(self, event):
        if self.radio.GetValue():
            self.choice = DIFFLIB
        elif self.radio2.GetValue():
            self.choice = PDML
        self.Destroy()

    def on_close(self, e):
        self.Destroy()

    @property
    def choice(self):
        return self._choice

    @choice.setter
    def choice(self, value):
        self._choice = value


class Smbcmp(wx.App):
    def OnInit(self):
        self.frame = MainFrame(None, wx.ID_ANY, "")
        self.SetTopWindow(self.frame)
        self.frame.Show()
        return True


if __name__ == "__main__":
    main()
