diff --git a/vpn_indicator.py b/vpn_indicator.py index f665fd3..bcd3a45 100755 --- a/vpn_indicator.py +++ b/vpn_indicator.py @@ -1,10 +1,10 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- -import select - __author__ = 'duc_tin' -from Queue import Empty +from Queue import Empty, Queue +from threading import Thread +import select import signal, os import socket, errno import time @@ -21,112 +21,167 @@ class InfoServer: def __init__(self, port): self.host = 'localhost' - self.data_payload = 2048 # buffer + self.port = port + self.buffer = 2048 # buffer self.backlog = 1 + + self.is_listening = False + self.is_connected = False self.client = None - self.state = False - self.port = port self.sock = socket.socket() self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) self.server_address = self.host, self.port + self.last_msg = '' - def wait_client(self): - self.client, address = self.sock.accept() - self.client.send('Hi') - self.client.sendall(self.last_msg) + self.readlist = [self.sock] # use for select - def run(self, q_in, q_out): - # print 'starting server at %s:%s' % server_address - self.sock.bind(self.server_address) - self.sock.listen(self.backlog) + def listen(self): + try: + self.sock.bind(self.server_address) + self.sock.listen(self.backlog) + print 'listening' - # wait forever for an incoming connection from indicator - self.client, address = self.sock.accept() # this line is not really necessary but just leave it there + return True + except socket.errno, e: + print e + return False + def check_io(self, q_info, q_cmd): + """Receive information about vpn tunnel + :type q_info: Queue + :type q_cmd: Queue + """ while True: + + # try to bind the socket + while not self.is_listening: + self.is_listening = self.listen() + time.sleep(2) + + # normal select protocol, timeout 0.5 sec + readable, _, _ = select.select(self.readlist, [], [], 0.5) + for s in readable: + if s is self.sock: # incoming connection + self.client, addrr = self.sock.accept() + self.readlist.append(self.client) + self.is_connected = True + print 'Server: Connected with %s:%s' % addrr + q_info.put('connected') + + else: # client sent something + try: + + data = self.client.recv(self.buffer) + if data: + print 'main sent: ', data + else: + self.is_connected = False + self.readlist.remove(self.client) + print 'main disconnected' + + q_info.put(data) + + except socket.error as e: + print 'Client die unexpectedly' + self.is_connected = False + + # get cmd from indicator without blocking try: - ready = select.select([self.client], [], [], 1) - if ready[0]: - data = self.client.recv(self.data_payload) - if 'close' in data: - self.client.close() - self.wait_client() - - info = q_in.get_nowait() - self.client.sendall(info) - self.last_msg = info + cmd = q_cmd.get_nowait() + if 'dead' in cmd: + print 'dead signal received' + self.sock.shutdown(socket.SHUT_RDWR) + self.sock.close() + return 0 + else: + self.send(cmd) except Empty: - time.sleep(1) - except socket.errno, e: - self.wait_client() - except IOError as e: - if e.errno == errno.EPIPE: - self.wait_client() - except Exception as e: - print 'unkown', str(e) + pass - def stop(self): - self.client.close() + def send(self, msg): + if self.is_connected: + self.client.sendall(msg) + return True + else: + return False class InfoClient: def __init__(self, port): - self.port = port self.host = 'localhost' - self.data_payload = 2048 # buffer + self.port = port + + self.buffer = 2048 # buffer self.sock = socket.socket() self.server_address = self.host, port - self.state = False + self.is_connected = False - def connect(self): - # print 'connect to %s:%s' % server_address - try: - self.sock = socket.create_connection(self.server_address) - print 'socket: connected' - return True - except socket.error, e: - print str(e) - - def get_data(self): - if not self.check_alive(): - self.state = False - return 'Offline' - elif not self.state: - self.state = True - return 'connected' - - data = '' - try: - ready = select.select([self.sock], [], [], 1) - if ready[0]: - data = self.sock.recv(self.data_payload) - except socket.errno, e: - print 'Socket error: ' + str(e) - except Exception, e: - print 'program error: '+ str(e) + self.last_msg = '' - return data + def connect(self): + while not self.is_connected: + try: + self.sock = socket.create_connection(self.server_address) + # print 'socket: connected' + self.is_connected = True + + # update current status + if self.last_msg: + self.send(self.last_msg) + + except socket.error, e: + # print str(e) + time.sleep(2) + + def check_io(self, q_info, q_cmd): + """Receive information about vpn tunnel + :type q_info: Queue, may not needed, we could send info directly + :type q_cmd: Queue + """ + while True: - def check_alive(self): - try: - self.sock.send('hello') - return True - except socket.error, e: - if 'Broken pipe' in e or 'Bad file descriptor' in e: - print 'server die' - self.sock.close() - return self.connect() + if self.is_connected: + # check if there is cmd from indicator + readable, _, _ = select.select([self.sock], [], [], 0.5) + try: + + data = self.sock.recv(self.buffer) + if data: + # print data + q_cmd.put(data) + else: + self.is_connected = False + # print 'server die' + + except socket.error as e: + print 'Server die unexpectedly' + self.is_connected = False else: - print 'Socket error: ' + str(e) + self.connect() + + def send(self, msg): + self.last_msg = msg + if self.is_connected: + try: + self.sock.sendall(msg) + return True + except socket.error: + return False + else: + return False class VPNIndicator: - def __init__(self, tcp_client): + def __init__(self, q_info, q_cmd): signal.signal(signal.SIGINT, self.handler) signal.signal(signal.SIGTERM, self.handler) - self.tcpClient = tcp_client + + # pipe for send/recv data to tcp server + self.q_info = q_info + self.q_cmd = q_cmd + self.APPINDICATOR_ID = 'myappindicator' self.icon1 = os.path.abspath('connected.svg') self.icon2 = os.path.abspath('connectnot.svg') @@ -146,7 +201,6 @@ def __init__(self, tcp_client): self.notifier.set_timeout(2) notify.init(self.APPINDICATOR_ID) - self.run() def run(self, *args): GLib.timeout_add(2000, self.callback, *args) @@ -170,16 +224,30 @@ def reload(self, data_in): self.indicator.set_icon(self.icon2) self.status('', ["Offline"]) self.hang = True + elif 'main exit' in data_in: + self.quit() return True def build_menu(self): menu = Gtk.Menu() - item_joke = Gtk.MenuItem('VPN Status') - item_joke.connect('activate', self.status, self.last_recv) - menu.append(item_joke) + # show status popup + current_status = Gtk.MenuItem('VPN Status') + current_status.connect('activate', self.status, self.last_recv) + menu.append(current_status) + + # connect to the next vpn on the list + next_vpn = Gtk.MenuItem('Next VPN') + next_vpn.connect('activate', self.send_cmd, 'next') + menu.append(next_vpn) + # connect to the next vpn on the list + stop_vpn = Gtk.MenuItem('Stop VPN') + stop_vpn.connect('activate', self.send_cmd, 'stop') + menu.append(stop_vpn) + + # quit button item_quit = Gtk.MenuItem('Quit') item_quit.connect('activate', self.handler, '') menu.append(item_quit) @@ -187,17 +255,20 @@ def build_menu(self): menu.show_all() return menu - def quit(self, source): + def quit(self, source=None): + # send dead signal to tcp server + self.q_cmd.put('dead') notify.uninit() Gtk.main_quit() - def status(self, _, messages=[0]): + def status(self, menu_obj, messages=''): """ :type messages: list """ - if _: - messages = self.last_recv + if not messages: + messages = ['unknown'] + self.notifier.close() if 'connected' in messages[0]: @@ -207,12 +278,12 @@ def status(self, _, messages=[0]): print messages[1:] summary = 'VPN tunnel established' body = ''' - %s \t %s - Ping: \t\t\t%s \tSpeed : \t%s Mbps - Up time:\t\t%s \tSeason: \t%s - Log: \t\t\t%s - Score: \t\t\t%s - Protocol: \t\t%s \tPortal: \t%s + %s \t %s + Ping: \t\t%s \tSpeed : %s Mbps + Up time:\t%s \tSeason: %s + Log: \t\t%s + Score: \t\t%s + Protocol: \t%s \tPortal: %s ''' % tuple(messages[1:]) elif 'terminate' in messages[0]: summary = 'VPN tunnel has broken' @@ -227,17 +298,30 @@ def status(self, _, messages=[0]): def handler(self, signal_num, frame): print 'interrupt now' self.quit('') + + def send_cmd(self, menu_obj, arg): + print 'indicator sent:', arg + self.q_cmd.put(arg) + + def callback(self): + try: - self.tcpClient.sock.send('close') - self.tcpClient.sock.close() - except socket.error: + data = self.q_info.get_nowait() + self.reload(data) + except Empty: pass - def callback(self): - data = self.tcpClient.get_data() - self.reload(data) return True + if __name__ == '__main__': - me = InfoClient(8088) - indicator = VPNIndicator(me) \ No newline at end of file + # queue for interacting between indicator and server + a, b = Queue(), Queue() + + server = InfoServer(8088) + t = Thread(target=server.check_io, args=(a, b)) # shouldn't be daemon + t.start() + + indicator = VPNIndicator(a, b) + indicator.run() + t.join() diff --git a/vpnproxy_tui.py b/vpnproxy_tui.py index 710d967..f239127 100755 --- a/vpnproxy_tui.py +++ b/vpnproxy_tui.py @@ -119,7 +119,7 @@ def __init__(self): self.vpn_server = None self.vpn_process = None self.vpn_queue = None - self.connect_status = False + self.is_connected = False self.kill = False self.get_limit = 1 @@ -180,7 +180,7 @@ def first_config(self): ip = socket.gethostbyname(proxy) if proxy: - print ' You are using proxy: '+ctext('%s:%s' % (proxy, port), 'pB') + print ' You are using proxy: '+ctext('%s:%s' % (proxy, port), 'bB') useit = 'yes' if raw_input( ctext(' Use this proxy? ', 'B') + '([yes]|no):') in 'yes' else 'no' @@ -322,10 +322,10 @@ def refresh_data(self, resort_only=''): self.messages['debug'].appendleft(' Sequence completed') def probe(self): - """ Filter out fetched dead Vpn Servers - """ + """ Filter out fetched dead Vpn Servers """ def is_alive(servers, queue): + """ Worker for threading""" target = [(self.vpndict[name].ip, self.vpndict[name].port) for name in servers] if self.use_proxy == 'yes': @@ -410,10 +410,8 @@ def vpn_output(out, queue): out.close() def vpn_connect(self, chosen): - """ - Disconnect the current connection and spawn a new one - """ - if self.connect_status: + """ Disconnect the current connection and spawn a new one """ + if self.is_connected: self.vpn_cleanup() server = self.vpndict[self.sorted[chosen]] @@ -439,7 +437,7 @@ def vpn_cleanup(self): if p.poll() is None: p.send_signal(signal.SIGINT) p.wait() - self.connect_status = False + self.is_connected = False self.dns_manager('restore') # make sure openvpn did close its device @@ -466,7 +464,7 @@ def vpn_checker(self): """ p, q = self.vpn_process, self.vpn_queue - if self.kill and self.connect_status: + if self.kill and self.is_connected: self.kill = False self.vpn_cleanup() self.messages['status'] += ['VPN tunnel is terminated', ''] @@ -482,8 +480,8 @@ def vpn_checker(self): self.dropped_time = 0 self.dns_manager('change') self.messages['status'] += ['VPN tunnel established successfully', 'Ctrl+C to quit VPN'] - self.connect_status = True - elif self.connect_status and 'Restart pause, ' in line and self.dropped_time <= self.max_retry: + self.is_connected = True + elif self.is_connected and 'Restart pause, ' in line and self.dropped_time <= self.max_retry: self.dropped_time += 1 self.messages['status'][1] = 'Vpn has restarted %s time(s)' % self.dropped_time elif 'Restart pause, ' in line or 'Cannot resolve' in line or 'Connection timed out' in line or 'SIGTERM' in line: @@ -496,7 +494,7 @@ def vpn_checker(self): elif '--http-proxy MUST' in line: self.messages['status'] += ['Can\'t use udp with proxy!', ' '] - elif p.poll() is None and not self.connect_status: + elif p.poll() is None and not self.is_connected: if 0 < self.dropped_time <= self.max_retry: self.messages['status'][0] = 'Connecting...' else: @@ -560,13 +558,14 @@ def __init__(self, vpn_connection): # indicator self.q2indicator = Queue() self.qfindicator = Queue() - - self.infoserver = InfoServer(8088) - self.indicator = Thread(target=self.infoserver.run, args=(self.q2indicator, self.qfindicator)) - self.indicator.daemon = True + + # should run on a thread so that it won't delay/block urwid + self.infoclient = InfoClient(8088) + self.indicator = Thread(target=self.infoclient.check_io, args=(self.q2indicator, self.qfindicator)) + self.indicator.daemon = True # client doesn't block port, it can die with main safely self.indicator.start() - self.sent = False - self.last_msg = '' + self.prev_status = False + # self.last_msg = '' def get_vpn_data(self): del self.data_ls[:] @@ -594,11 +593,11 @@ def periodic_checker(self, loop, user_data=None): if self.clear_input: self.input.set_edit_text(self.clear_input[1]) self.clear_input = False - - if self.ovpn.connect_status != self.sent: - self.sent = self.ovpn.connect_status - self.communicator((self.sent, '')) - + + # send/recv information to/from indicator + self.communicator() + + # refresh the terminal screen if self.cache_msg != self.ovpn.messages: self.status(self.ovpn.messages) self.cache_msg = deepcopy(self.ovpn.messages) @@ -609,7 +608,7 @@ def periodic_checker(self, loop, user_data=None): def signal_handler(self, signum, frame): self.ovpn.kill = True self.printf("Ctrl C is pressed. Press again or 'q' to quit program") - if not self.ovpn.connect_status: + if not self.ovpn.is_connected: raise urwid.ExitMainLoop() def input_handler(self, Edit, key_ls=None): @@ -900,12 +899,22 @@ def status(self, msg=None): ind += 1 - def communicator(self, (msg, ovpn)): - if msg: - msgs = 'successfully;' + repr(self.ovpn.vpn_server) - self.q2indicator.put(msgs) - else: - self.q2indicator.put('terminate') + def communicator(self): + # send info + if self.ovpn.is_connected != self.prev_status: + self.prev_status = self.ovpn.is_connected + if self.prev_status: + msgs = 'successfully;' + repr(self.ovpn.vpn_server) + self.infoclient.send(msgs) + else: + self.infoclient.send('terminate') + + # receive cmd + try: + cmd = self.qfindicator.get_nowait() + self.printf('Indicator told: ' + cmd) + except Empty: + pass def run(self): self.loop.run()