#!/usr/bin/python
#
# Copyright (C) 2010-2012 Red Hat, Inc.
# Authors:
# Thomas Woerner <twoerner@redhat.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 2 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/>.
#

from gi.repository import Gtk, GLib, GObject, Gio, Notify, NMClient

import os
import dbus.mainloop.glib
import slip.dbus
import types

from firewall.config import *
from firewall.config.dbus import *
from firewall.client import FirewallClient
import dbus


PATH = [ ]
for p in os.getenv("PATH").split(":"):
    if p not in PATH:
        PATH.append(p)

def search_app(app):
    for p in PATH:
        _app = "%s/%s" % (p, app)
        if os.path.exists(_app):
            return _app
    return None

APPLET_SCHEMA = "org.fedoraproject.FirewallApplet"

class TrayApplet(object):
    @staticmethod
    def position_function(menu, icon):
        return (Gtk.StatusIcon.position_menu(menu, icon))

    def __init__(self, use_dbus=True):
        self.name = _("Firewall Applet")
#        self.comment = _("")
        self.icon_name = "firewall-applet"
        self.settings = Gio.Settings.new(APPLET_SCHEMA)

        self.icons = { "normal": None, "error": None, "panic": None, }
        self.timer = None
        self.mode = None
        self._blink = False
        self.blink_count = 0

        self.active_zones = { }
        self.connections = { }

        self.statusicon = Gtk.StatusIcon.new()
#        self.statusicon.connect("delete_event", Gtk.main_quit)
        self.statusicon.set_from_icon_name(self.icon_name)

        theme = Gtk.IconTheme.get_default()
        # Gtk.IconSize.MENU
        size = 24

        info = theme.lookup_icon("firewall-applet", size, 0)
        if not info:
            print("Icon 'firewall-applet' could not be loaded")
        else:
            self.icons["normal"] = info.load_icon()

        for _type in [ "error", "panic" ]:
            info = theme.lookup_icon("firewall-applet-%s" % _type, size, 0)
            if not info:
                print("Icon 'firewall-applet-%s' could not be loaded" % _type)
            else:
                self.icons[_type] = info.load_icon()

        self.right_menu = Gtk.Menu.new()

#        item = Gtk.CheckMenuItem.new_with_mnemonic(_("Enable Firewall"))
#        item.set_active(True)
#        item.set_sensitive(False)
#        self.right_menu.append(item)

        self.notification_check = Gtk.CheckMenuItem.new_with_mnemonic(
            _("Enable Notifications"))
        self.notification_check.set_active(self.settings.get_boolean(
                "notifications"))
        self.settings.connect("changed::notifications",
                              self.settings_check_changed,
                              self.notification_check)
        self.notification_check.connect('toggled', self.check_toggled,
                                        self.settings, "notifications")
        self.right_menu.append(self.notification_check)

        self.right_menu.append(Gtk.SeparatorMenuItem.new())

        item = Gtk.MenuItem.new_with_mnemonic(_("Edit Firewall Settings..."))
        item.connect("activate", self.configure_cb)
        if not search_app("firewall-config"):
            item.set_sensitive(False)

        self.right_menu.append(item)
        
        self.right_menu.append(Gtk.SeparatorMenuItem.new())

        self.panic_check = Gtk.CheckMenuItem.new_with_mnemonic(\
            _("Block all network traffic"))
        self.panic_check.set_active(False)
        self.panic_check_id = self.panic_check.connect("toggled",
                                                       self.panic_mode_cb)
        self.right_menu.append(self.panic_check)

        self.right_menu.append(Gtk.SeparatorMenuItem.new())

        item = Gtk.ImageMenuItem.new_from_stock(Gtk.STOCK_ABOUT, None)
        item.connect('activate', self.about_cb)
        self.right_menu.append(item)

        self.statusicon.connect("popup-menu", self.right_menu_cb,
                                self.right_menu)

        self.about_dialog = Gtk.AboutDialog.new()
        self.about_dialog.set_name(self.name)
        self.about_dialog.set_version(VERSION)
#        self.about_dialog.set_comments(self.comment)
        self.about_dialog.set_license(LICENSE)
        self.about_dialog.set_wrap_license(True)
        self.about_dialog.set_copyright(COPYRIGHT)
        self.about_dialog.set_authors(AUTHORS)
        self.about_dialog.set_logo_icon_name(self.icon_name)

        self.statusicon.set_visible(True)

        dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)
        try:
            self.bus = slip.dbus.SystemBus()
            self.bus.default_timeout = None
        except Exception as msg:
            print(_("Not using slip"), msg)
            self.bus = dbus.SystemBus()

        self.bus.add_signal_receiver(
            handler_function=self.dbus_connection_changed,
            signal_name="NameOwnerChanged",
            dbus_interface="org.freedesktop.DBus")

        self.bus.add_signal_receiver(
            self.signal_receiver,
            dbus_interface=DBUS_INTERFACE_ZONE,
            interface_keyword='interface',
            member_keyword='member',
            path_keyword='path')

        self.bus.add_signal_receiver(
            self.signal_receiver,
            dbus_interface=DBUS_INTERFACE,
            interface_keyword='interface',
            member_keyword='member',
            path_keyword='path')

        self.bus.add_signal_receiver(
            self.nm_signal_receiver,
            dbus_interface="org.freedesktop.NetworkManager",
            sender_keyword='sender')

        self.fw = None
        self.nm_signal_receiver()
        self.connect_to_firewalld()

        self.update()

    def settings_check_changed(self, settings, key, button):
        button.set_active(settings.get_boolean(key))

    def check_toggled(self, button, settings, key):
        settings.set_boolean(key, button.get_active())

    def notify(self, msg, sender=None, urgency=Notify.Urgency.NORMAL):
        n = Notify.Notification.new(self.name, msg, self.icon_name)
        n.set_urgency(urgency)
        try:
            n.show()
        except:
            pass

    def dbus_connection_changed(self, name, old_owner, new_owner):
        if name != DBUS_INTERFACE:
            return

        if new_owner:
            # new connection
            self.connect_to_firewalld()
        else:
            # lost connection
            self.connection_to_firewalld_lost()

        # no change, no notification
        if new_owner and not self.fw:
            return
        if not new_owner and self.fw:
            return

        ed = { 1: _("Connection to FirewallD lost."),
               0: _("Connection to FirewallD established.") }
        self.notify(ed[self.fw == None], urgency=Notify.Urgency.CRITICAL)

    def connect_to_firewalld(self):
        try:
            self.fw = FirewallClient(self.bus)
        except dbus.DBusException as e:
            if e._dbus_error_name == \
                    'org.freedesktop.DBus.Error.ServiceUnknown':
                # no connection
                return
            else:
                raise

        self.set_mode("normal")
        self.active_zones.clear()
        active_zones = self.fw.getActiveZones()
        if active_zones:
            self.active_zones = active_zones
        self.update()

    def connection_to_firewalld_lost(self):
        if self.fw:
            self.fw = None
            self.active_zones.clear()
            self.set_mode("error")
        self.update()

    def panic_mode_cb(self, check):
        if not self.fw:
            return
        if check.get_active():
            self.fw.enablePanicMode()
        else:
            self.fw.disablePanicMode()
        
    def right_menu_cb(self, widget, button, time, menu):
        if button != 3:
            return
        menu.show_all()
        menu.popup(None, None, self.position_function,
                   self.statusicon, button, time)

    def about_cb(self, widget):
        self.about_dialog.run()
        self.about_dialog.hide()

    def configure_cb(self, widget):
        os.system("firewall-config &")

    def __blink(self, arg=None):
        if self.blink_count != 0:
            if self.blink_count > 0 and self._blink:
                self.blink_count -= 1
            self._blink = not self._blink
            self.timer = GLib.timeout_add_seconds(1, self.__blink, None)

        if not self._blink:
            self.statusicon.set_from_pixbuf(self.icons[self.mode])
        else:
            self.statusicon.set_from_pixbuf(self.icons["normal"])

    def set_mode(self, mode, count=5):
        if self.mode != mode:
            if self.timer:
                GLib.source_remove(self.timer)
                self.timer = None
                self._blink = False
            self.mode = mode

        elif self.mode == mode and self.timer:
            if self.blink_count == 0:
                self.blink_count += 1
            return

        if mode == "normal":
            self.statusicon.set_from_pixbuf(self.icons[mode])
            return

        if count != 0:
            self._blink = True
            self.blink_count = count
            self.__blink()

    def update(self):
        if not self.fw:
            self.set_mode("error")
            self.tooltip = "<span color='#FF0000'>" + \
                _("No connection to firewall daemon") + "</span>"
            self.statusicon.set_tooltip_markup(self.tooltip)
            return

        if self.panic_check.get_active():
            self.tooltip = "<big><b><span color='#FF0000'>" + \
                _("PANIC MODE") + "</span></b></big>"
            self.statusicon.set_tooltip_markup(self.tooltip)
            return
        
        self.tooltip = ""

        messages = [ ]
        if len(self.active_zones) > 0:
            for zone in sorted(self.active_zones):
                for interface in sorted(self.active_zones[zone]):
                    if interface in self.connections:
                        connection = self.connections[interface]
                        text = _("Zone '{zone}' active for connection "
                                 "'{connection}' on interface '{interface}'")
                    else:
                        text = _("Zone '{zone}' active for interface "
                                 "'{interface}'")
                        connection = None
                    messages.append(text.format(zone=zone,
                                                connection=connection,
                                                interface=interface))
            self.tooltip = "\n".join(messages)
        else:
            self.tooltip = _("No Active Zones.")

        self.statusicon.set_tooltip_markup(self.tooltip)

    def signal_receiver(self, *args, **kwargs):
#        print(args, kwargs)
        if not "member" in kwargs:
            return
        signal = kwargs["member"]

        if self.notification_check.get_active():
            if signal == "Reloaded":
                self.notify(_("FirewallD has been reloaded."))
            elif signal == "DefaultZoneChanged":
                self.notify(_("Default zone changed to '%s'.") % args[0])

        if signal in [ "PanicModeEnabled", "PanicModeDisabled" ]:
            if signal == "PanicModeEnabled":
                enable = True
            else:
                enable = False

            self.panic_check.handler_block(self.panic_check_id)
            self.panic_check.set_active(enable == 1)
            self.panic_check.handler_unblock(self.panic_check_id)

            self.update()

            if enable:
                self.set_mode("panic")
            else:
                self.set_mode("normal")

            if self.notification_check.get_active():
                ed = { 1: _("All network traffic is blocked."),
                       0: _("Network traffic is not blocked anymore.") }
                self.notify(ed[enable], urgency=Notify.Urgency.CRITICAL)

        elif signal in [ "InterfaceAdded", "InterfaceRemoved" ]:
            if signal == "InterfaceAdded":
                enable = True
            else:
                enable = False

            # interface signal
            zone = str(args[0])
            interface = str(args[1])

            if enable:
                if zone not in self.active_zones:
                    self.active_zones[zone] = [ ]
                if interface not in self.active_zones[zone]:
                    self.active_zones[zone].append(interface)
            else:
                try:
                    self.active_zones[zone].remove(interface)
                    if len(self.active_zones[zone]) < 1:
                        del self.active_zones[zone]
                except Exception as msg:
                    print(msg)

            self.update()

            # send notification if enabled
            if self.notification_check.get_active():
                ed = { 1: _("activated"),
                       0: _("deactivated") }
                if interface in self.connections:
                    connection = self.connections[interface]
                    text = _("Zone '{zone}' {activated_deactivated} for "
                             "connection '{connection}' on "
                             "interface '{interface}'")
                else:
                    connection = None
                    text = _("Zone '{zone}' {activated_deactivated} for "
                             "interface '{interface}'")
                self.notify(text.format(zone=zone,
                                        activated_deactivated=ed[enable],
                                        connection=connection,
                                        interface=interface))

        elif signal == "ZoneChanged":
            zone = args[0]
            interface = args[1]

            for _zone in self.active_zones:
                if interface in self.active_zones[zone]:
                    self.active_zones[zone].remove(str(interface))
                    if len(self.active_zones[zone]) < 1:
                        del self.active_zones[zone]
                        break

            self.update()

            self.active_zones.setdefault(zone, [ ]).append(str(interface))
            if self.notification_check.get_active():
                self.notify(_("Zone '%s' activated for interface '%s'") % \
                                (zone, interface))


    def nm_signal_receiver(self, *args, **kwargs):
#        print(args, kwargs)
        self.connections.clear()
        try:
            self.nmclient = NMClient.Client(\
                dbus_path='/org/freedesktop/NetworkManager')
            active = self.nmclient.get_active_connections()
            for a in active:
                obj = self.bus.get_object("org.freedesktop.NetworkManager",
                                          a.get_connection())
                connection = obj.GetSettings()['connection']['id']
                for dev in a.get_devices():
                    self.connections[dev.get_iface()] = connection
        except Exception as msg:
            print(msg)

        self.update()

# MAIN

if __name__ == "__main__":
    mainloop = GObject.MainLoop()
    applet = TrayApplet()
    Notify.init(applet.name)
    mainloop.run()
