#!/usr/bin/env python

# ©2020 Shzublox, BrainreckedTech
# Licensed under the GPL v3

import subprocess, wx, wx.adv

PROGRAM_NAME = 'SSH Proxy'
STOPPED_ICON = "/usr/share/pixmaps/sshproxy-unused.png"
STARTED_ICON = "/usr/share/pixmaps/sshproxy-active.png"
DEFAULT_PORT = "8080"
DEFAULT_ADDR = "proxy"

def create_menu_item(menu, label, func):
    item = wx.MenuItem(menu, -1, label)
    menu.Bind(wx.EVT_MENU, func, id=item.GetId())
    menu.Append(item)
    return item

class MyTaskBarIcon(wx.adv.TaskBarIcon):
    def __init__(self, frame):
        wx.adv.TaskBarIcon.__init__(self)
        self.frame = frame

        self.SetIcon(self.frame.stopped_icon, PROGRAM_NAME)
        self.Bind(wx.adv.EVT_TASKBAR_LEFT_DOWN, self.leftclicked)

    def CreatePopupMenu(self):
        menu = wx.Menu()
        create_menu_item(menu, 'Configure', self.configure)
        create_menu_item(menu, 'Exit', self.frame.close)
        return menu

    def leftclicked(self, event):
        if self.frame.ssh:
            self.frame.on_stop(None)
        else:
            self.frame.on_start(None)
    
    def configure(self, event):
        self.frame.Show()
        self.frame.Restore()

class MyFrame(wx.Frame):
    def __init__(self, parent, stopped_icon, started_icon):
        wx.Frame.__init__(self, parent, title=PROGRAM_NAME,
                          style = wx.DEFAULT_FRAME_STYLE &
                          ~(wx.RESIZE_BORDER | wx.MINIMIZE_BOX
                            | wx.MAXIMIZE_BOX | wx.CLOSE_BOX))

        self.ssh = None
        self.stopped_icon = stopped_icon
        self.started_icon = started_icon

        self.SetIcon(stopped_icon)

        self.tray = MyTaskBarIcon(self)
        self.Bind(wx.EVT_ICONIZE, self.iconize)
        self.Bind(wx.EVT_CLOSE, self.close)

        panel_sizer = wx.FlexGridSizer(cols=1)
        text_sizer = wx.FlexGridSizer(cols=2)
        button_sizer = wx.BoxSizer()
        panel = wx.Panel(self)

        portlabel = wx.StaticText(panel, label="Port: ")
        self.portbox = wx.TextCtrl(panel, value=DEFAULT_PORT, size=(300, -1))

        remotelabel = wx.StaticText(panel, label="Remote: ")
        self.remotebox = wx.TextCtrl(panel, value=DEFAULT_ADDR, size=(300, -1))

        self.startbutton = wx.Button(panel, label="Start")
        self.stopbutton = wx.Button(panel, label="Stop")

        self.stopbutton.Disable()

        text_sizer.Add(portlabel, 1, wx.EXPAND)
        text_sizer.Add(self.portbox, 1, wx.EXPAND)
        text_sizer.Add(remotelabel, 1, wx.EXPAND)
        text_sizer.Add(self.remotebox, 1, wx.EXPAND)
        text_sizer.Fit(self)

        button_sizer.Add(self.startbutton, 1, wx.EXPAND)
        button_sizer.Add(self.stopbutton, 1, wx.EXPAND)
        button_sizer.Fit(self)

        panel_sizer.Add(text_sizer, 1, wx.EXPAND)
        panel_sizer.Add(button_sizer, 1, wx.EXPAND)
        panel_sizer.Fit(self)

        self.startbutton.Bind(wx.EVT_BUTTON, self.on_start)
        self.stopbutton.Bind(wx.EVT_BUTTON, self.on_stop)

        self.SetSizer(panel_sizer)

    def close(self, event):
        self.tray.RemoveIcon()
        self.tray.Destroy()
        self.Destroy()

    def iconize(self, event):
        if self.IsIconized():
            self.Show(False)
            self.Restore()

    def on_start(self, event):
        if self.portbox.GetValue() == "" or self.remotebox.GetValue() == "":
            return

        self.startbutton.Disable()
        self.stopbutton.Enable()
        self.tray.SetIcon(self.started_icon)
        self.ssh = subprocess.Popen(["ssh", "-ND",
                                     "%s" % self.portbox.GetValue(),
                                     "%s" % self.remotebox.GetValue()],
                                    stdin = subprocess.PIPE,
                                    stdout = subprocess.PIPE,
                                    stderr = subprocess.PIPE,
                                    universal_newlines = True)

    def on_stop(self, event):
        self.startbutton.Enable()
        self.stopbutton.Disable()
        self.tray.SetIcon(self.stopped_icon)
        if self.ssh:
            self.ssh.terminate()
            self.ssh = None

def main():
    app = wx.App(False)

    image = wx.Image(STOPPED_ICON, wx.BITMAP_TYPE_ANY)
    bmp = wx.Bitmap(image)
    stopped_icon = wx.Icon()
    stopped_icon.CopyFromBitmap(bmp)

    image = wx.Image(STARTED_ICON, wx.BITMAP_TYPE_ANY)
    bmp = wx.Bitmap(image)
    started_icon = wx.Icon()
    started_icon.CopyFromBitmap(bmp)

    frame = MyFrame(None, stopped_icon, started_icon)
    app.MainLoop()

if __name__ == "__main__":
    main()