<repo_name>strangelarry/HostsQuick<file_sep>/hostsview.py
from gi.repository import Gtk


class HostsTextView(Gtk.TextView):
    def __init__(self):
        super(Gtk.TextView, self).__init__()
        self.buffer_list = {}

    def add_to_buffer(self, hosts_name, editable):
        buffer = Gtk.TextBuffer()
        self.buffer_list[hosts_name] = (buffer, editable)

    def remove_from_buffer(self, hosts_name):
        if hosts_name in self.buffer_list:
            del self.buffer_list[hosts_name]

    def set_buffer_content(self, hosts_name, content):
        buffer, editable = self.buffer_list.get(hosts_name)
        if buffer is not None:
            buffer.set_text(content)

    def get_buffer_content(self, hosts_name):
        buffer, editable = self.buffer_list.get(hosts_name)
        if buffer is not None:
            content = buffer.get_text(*buffer.get_bounds(), True)
            return content
        return None

    def set_buffer(self, hosts_name):
        buffer, editable = self.buffer_list.get(hosts_name)
        if buffer is None:
            return
        super().set_buffer(buffer)
        super().set_editable(editable)

    def get_buffer(self):
        return super().get_buffer()

    def has_hosts(self, hosts_name):
        return hosts_name in self.buffer_list


class HostsListStore(Gtk.ListStore):
    def __init__(self, *args):
        super().__init__(*args)
        self.current_treeiter = None

    def remove(self, treeiter):
        if treeiter is None:
            return
        super().remove(treeiter)
        if len(self) == 0:
            self.current_treeiter = None


class HostsInputDialog(Gtk.Dialog):
    def __init__(self, parent, title):
        Gtk.Dialog.__init__(self, title, parent, 0,
                            (Gtk.STOCK_CANCEL, Gtk.ResponseType.CANCEL,
                             Gtk.STOCK_OK, Gtk.ResponseType.OK))
        self.entry = Gtk.Entry()
        box = self.get_content_area()
        box.add(self.entry)
        self.show_all()

    def get_entry(self):
        return self.entry.get_text()

    def set_entry_visible(self, visible):
        self.entry.set_visibility(visible)
<file_sep>/button_event.py
from hostsview import HostsInputDialog
from gi.repository import Gtk
from loader import HostsLoader
import sys


def on_menu_choose_hosts(obj, window, hosts_name):
    if not obj.get_active():
        return
    window.choose_type = 'menu'
    choose_hosts(window, hosts_name)


def on_window_choose_hosts(obj, window):
    window.choose_type = 'window'
    model = window.host_store
    treeiter = model.current_treeiter
    hosts_name = ''
    if treeiter is not None:
        hosts_name = model[treeiter][0]
    choose_hosts(window, hosts_name)


def choose_hosts(window, hosts_name=''):
    if window.choose_type is None or hosts_name == '':
        return
    current_hosts = HostsLoader.get_current_hosts()
    if current_hosts == hosts_name:
        return

    if HostsLoader.set_hosts(hosts_name):
        column = window.host_tree.get_column(0)
        column.set_title('Hosts List---->' + hosts_name)
        current_hosts = HostsLoader.get_current_hosts()
        set_menu_hosts(window, current_hosts)
        return
    dialog = HostsInputDialog(window, 'input root password')
    dialog.set_entry_visible(False)
    response = dialog.run()
    passwd = dialog.get_entry()

    if response == Gtk.ResponseType.OK:
        HostsLoader.set_sudo_passwd(passwd)
        if not HostsLoader.set_hosts(hosts_name):
            warn = Gtk.MessageDialog(window, 0, Gtk.MessageType.ERROR,
                                     Gtk.ButtonsType.CANCEL, '密码不正确！')
            warn.run()
            warn.destroy()
        else:
            column = window.host_tree.get_column(0)
            column.set_title('Hosts List---->' + hosts_name)
    elif response == Gtk.ResponseType.CANCEL:
        pass

    current_hosts = HostsLoader.get_current_hosts()
    set_menu_hosts(window, current_hosts)

    dialog.destroy()


def create_hosts(obj, window):
    model = window.host_store
    textview = window.textview

    dialog = HostsInputDialog(window, 'input hosts name')
    response = dialog.run()
    hosts_name = dialog.get_entry()

    if response == Gtk.ResponseType.OK:
        if hosts_name == '' or textview.has_hosts(hosts_name):
            warn = Gtk.MessageDialog(window, 0, Gtk.MessageType.ERROR,
                                     Gtk.ButtonsType.CANCEL, '输入为空或者hosts已存在！')
            warn.run()
            warn.destroy()
        else:
            model.append([hosts_name])
            textview.add_to_buffer(hosts_name, True)
            HostsLoader.add_hosts(hosts_name)

            group = window.menu.get_children()[0].get_group()
            hosts_item = Gtk.RadioMenuItem.new_with_label(group, hosts_name)
            window.menu.insert(hosts_item, len(window.menu.get_children())-2)
            hosts_item.connect('activate', on_menu_choose_hosts, window, hosts_name)
            hosts_item.show()
    elif response == Gtk.ResponseType.CANCEL:
        pass

    dialog.destroy()


def delete_hosts(obj, window):
    model = window.host_store
    textview = window.textview

    treeiter = model.current_treeiter
    if treeiter is not None:
        hosts_name = model[treeiter][0]
        if hosts_name == HostsLoader.SYS_HOSTS:
            warn = Gtk.MessageDialog(window, 0, Gtk.MessageType.ERROR,
                                     Gtk.ButtonsType.OK, '系统默认hosts不能删除！')
            warn.run()
            warn.destroy()
            return
        textview.remove_from_buffer(hosts_name)
        model.remove(treeiter)
        for item in window.menu.get_children():
            if hosts_name == item.get_label():
                window.menu.remove(item)
                window.menu.show_all()
                break
        HostsLoader.delete_hosts(hosts_name)


def save_content(obj, window):
    model = window.host_store
    textview = window.textview

    treeiter = model.current_treeiter
    if treeiter is not None:
        hosts_name = model[treeiter][0]
        if hosts_name == HostsLoader.SYS_HOSTS:
            return
        content = textview.get_buffer_content(hosts_name)
        HostsLoader.write_hosts(hosts_name, content)


def clear_content(obj, window):
    model = window.host_store
    textview = window.textview

    treeiter = model.current_treeiter
    if treeiter is not None:
        hosts_name = model[treeiter][0]
        if hosts_name == HostsLoader.SYS_HOSTS:
            return
        textview.set_buffer_content(hosts_name, '')


def reload_content(obj, window):
    model = window.host_store
    textview = window.textview

    treeiter = model.current_treeiter
    if treeiter is not None:
        hosts_name = model[treeiter][0]
        if hosts_name == HostsLoader.SYS_HOSTS:
            return
        content = HostsLoader.read_hosts(hosts_name)
        textview.set_buffer_content(hosts_name, content)


def on_host_tree_selection_changed(obj, window):
    model = window.host_store
    textview = window.textview

    _, treeiter = obj.get_selected()
    if treeiter is not None:
        model.current_treeiter = treeiter
        hosts_name = model[treeiter][0]
        textview.set_buffer(hosts_name)


def set_menu_hosts(window, hosts_name):
    for item in window.menu.get_children():
        if hosts_name == item.get_label():
            if not item.get_active():
                item.set_active(True)
            break


def quit_app(obj):
    sys.exit(0)
<file_sep>/loader.py
import os
import glob


class HostsLoader(object):
    HOSTS_FOLDER = './hosts'
    HOSTS_FILES = []
    SYS_HOSTS = 'sys'
    SUDO_PASSWD = ''
    HOSTS_REAL = '/etc/hosts'

    @classmethod
    def init(cls):
        if not os.path.exists(cls.HOSTS_FOLDER):
            os.mkdir(cls.HOSTS_FOLDER)
        if not cls.HOSTS_FILES:
            cls.HOSTS_FILES = glob.glob(cls.HOSTS_FOLDER + '/hosts.*')
        sys_hosts = cls.get_hosts_file(cls.SYS_HOSTS)
        if sys_hosts in cls.HOSTS_FILES:
            cls.HOSTS_FILES.remove(sys_hosts)
        else:
            with open('/etc/hosts') as f:
                content = f.read()
                cls.write_hosts(cls.SYS_HOSTS, content)
        cls.HOSTS_FILES.insert(0, sys_hosts)

    @classmethod
    def get_hosts_file(cls, hosts_name):
        return cls.HOSTS_FOLDER + '/hosts.' + hosts_name

    @classmethod
    def add_hosts(cls, hosts_name):
        hosts_file = cls.get_hosts_file(hosts_name)
        cls.HOSTS_FILES.append(hosts_file)
        cls.write_hosts(hosts_name)

    @classmethod
    def write_hosts(cls, hosts_name, content=''):
        hosts_file = cls.get_hosts_file(hosts_name)
        with open(hosts_file, 'w+') as f:
            f.write(content)

    @classmethod
    def read_hosts(cls, hosts_name):
        hosts_file = cls.get_hosts_file(hosts_name)
        with open(hosts_file) as f:
            content = f.read()
        return content

    @classmethod
    def delete_hosts(cls, hosts_name):
        hosts_file = cls.get_hosts_file(hosts_name)
        try:
            os.remove(hosts_file)
        except:
            pass

    @classmethod
    def load_hosts(cls):
        return [os.path.basename(hosts_name).split('.')[1] for hosts_name in cls.HOSTS_FILES]

    @classmethod
    def set_sudo_passwd(cls, passwd):
        cls.SUDO_PASSWD = <PASSWORD>

    @classmethod
    def set_hosts(cls, hosts_name):
        hosts_file = cls.get_hosts_file(hosts_name)

        result = os.system('echo %s|sudo -S ln -sf %s /etc/hosts'
                           % (cls.SUDO_PASSWD, os.path.abspath(hosts_file)))

        return not bool(result)

    @classmethod
    def get_current_hosts(cls):
        realpath = os.path.realpath(cls.HOSTS_REAL)
        try:
            hosts_name = os.path.basename(realpath).split('.')[1]
            return hosts_name
        except:
            return None
<file_sep>/main.py
import gi
gi.require_version('Gtk', '3.0')
gi.require_version('AppIndicator3', '0.1')
from gi.repository import Gtk, AppIndicator3
from button_event import *
from hostsview import HostsTextView, HostsListStore
from loader import HostsLoader


class MainWindow(Gtk.ApplicationWindow):

    def __init__(self, app):
        Gtk.Window.__init__(self, title='HostsQuick', application=app)
        self.set_default_size(800, 600)
        self.set_border_width(10)

        self.appindicator_id = 'HostsQuick'
        self.indicator = AppIndicator3.Indicator.new(self.appindicator_id, 'whatever',
                                                     AppIndicator3.IndicatorCategory.SYSTEM_SERVICES)
        self.menu = Gtk.Menu()

        self.outer_box = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=10)
        self.add(self.outer_box)

        self.left_box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6)
        self.left_box.set_size_request(250, 600)
        self.host_store = HostsListStore(str)
        self.host_tree = Gtk.TreeView(self.host_store)

        self.right_box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6)
        self.right_box.set_size_request(550, 600)
        self.textview = HostsTextView()

        self.outer_box.pack_start(self.left_box, True, True, 0)
        self.outer_box.pack_start(self.right_box, True, True, 0)

        self.choose_type = None

        HostsLoader.init()
        self.init_hosts()
        self.create_leftbox()
        self.create_rightbox()
        self.create_indicator()
        self.init_others()

        self.set_position(Gtk.WindowPosition.CENTER)

    def create_leftbox(self):
        renderer = Gtk.CellRendererText()
        column = Gtk.TreeViewColumn("Hosts List", renderer, text=0)
        self.host_tree.append_column(column)
        self.left_box.pack_start(self.host_tree, True, True, 0)
        select = self.host_tree.get_selection()
        select.connect("changed", on_host_tree_selection_changed, self)

        button_choose = Gtk.Button(label='Choose')
        button_choose.connect('clicked', on_window_choose_hosts, self)
        button_create = Gtk.Button(label='Create')
        button_create.connect('clicked', create_hosts, self)
        button_delete = Gtk.Button(label='Delete')
        button_delete.connect('clicked', delete_hosts, self)
        vbox = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL)
        vbox.pack_start(button_choose, False, False, 0)
        vbox.pack_start(button_create, False, False, 0)
        vbox.pack_start(button_delete, False, False, 0)
        self.left_box.pack_start(vbox, False, True, 0)

    def create_rightbox(self):
        scrolledwindow = Gtk.ScrolledWindow()
        scrolledwindow.set_hexpand(True)
        scrolledwindow.set_vexpand(True)
        scrolledwindow.add(self.textview)
        self.right_box.pack_start(scrolledwindow, True, True, 0)

        button_save = Gtk.Button(label='Save')
        button_save.connect('clicked', save_content, self)
        button_clear = Gtk.Button(label='Clear')
        button_clear.connect('clicked', clear_content, self)
        button_reload = Gtk.Button(label='Reload')
        button_reload.connect('clicked', reload_content, self)
        vbox = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL)
        vbox.pack_end(button_reload, False, False, 0)
        vbox.pack_end(button_clear, False, False, 0)
        vbox.pack_end(button_save, False, False, 0)
        self.right_box.pack_start(vbox, False, True, 0)

    def init_hosts(self):
        hosts = HostsLoader.load_hosts()
        for host in hosts:
            self.host_store.append([host])
            if host == HostsLoader.SYS_HOSTS:
                self.textview.add_to_buffer(host, False)
            else:
                self.textview.add_to_buffer(host, True)
            content = HostsLoader.read_hosts(host)
            self.textview.set_buffer_content(host, content)

    def init_others(self):
        current_hosts = HostsLoader.get_current_hosts()
        if current_hosts is not None:
            self.host_tree.get_column(0).set_title('Hosts List---->' + current_hosts)

    def create_indicator(self):
        self.indicator.set_status(AppIndicator3.IndicatorStatus.ACTIVE)
        group = []
        current_hosts = HostsLoader.get_current_hosts()

        for hosts_name in HostsLoader.load_hosts():
            hosts_item = Gtk.RadioMenuItem.new_with_label(group, hosts_name)
            group = hosts_item.get_group()
            self.menu.append(hosts_item)
            if hosts_name == current_hosts:
                hosts_item.set_active(True)
            hosts_item.connect('activate', on_menu_choose_hosts, self, hosts_name)

        separator = Gtk.SeparatorMenuItem()
        item_quit = Gtk.MenuItem('Quit')
        item_quit.connect('activate', quit_app)

        self.menu.append(separator)
        self.menu.append(item_quit)
        self.menu.show_all()
        self.indicator.set_menu(self.menu)


class MyApplication(Gtk.Application):
    def __init__(self):
        Gtk.Application.__init__(self)

    def do_activate(self):
        win = MainWindow(self)
        win.show_all()

    def do_startup(self):
        Gtk.Application.do_startup(self)


if __name__ == '__main__':
    app = MyApplication()
    exit_status = app.run(sys.argv)
    sys.exit(exit_status)
