Merge pull request #1 from Akiiino/master

Add small changes, make code cleaner

teletun.py

@@ -16,42 +16,37 @@ import signal
 
 
 # Connect to telegram
-print 'Connecting to telegram...'
+print('Connecting to telegram...')
 
 receiver = Receiver(host="localhost", port=4458)
 sender = Sender(host="localhost", port=4458)
 
 # Retrieve contact list
-contacts = {}
 
 try:
-    counter = 0
-    for c in sender.dialog_list():
-        counter += 1
-        contacts[counter] = c
-        print unicode(counter) + ': \t' + unicode(c['print_name'])
+    contacts = [c for c in sender.dialog_list()]
+    for i, user in enumerate(contacts):
+        print(unicode(i) + ': \t' + unicode(user['print_name']))
 except ConnectionError:
-    print 'Could not connect to telegram-cli. Start it by issuing "telegram-cli --json -P 4458" in a separate console.'
-    exit(1)
+    print('Could not connect to telegram-cli. Start it by issuing "telegram-cli --json -P 4458" in a separate console.')
+    sys.exit(1)
 
 # Ask user to choose contact
-i = input('Telegram online, please enter contact to connect to (by number): ')
-
-if i not in contacts:
-    print 'Please enter a number in the above range!'
-    exit(1)
-
+i = int(input('Telegram online, please enter contact to connect to (by number): '))
 
 # Print username
-username = unicode(contacts[i]['print_name'])
-peer_id = contacts[i]['peer_id']
-print 'Connecting to partner: ' + username
-
+try:
+    username = unicode(contacts[i]['print_name'])
+    peer_id = contacts[i]['peer_id']
+    print('Connecting to partner: ' + username)
+except IndexError:
+    print('Please enter a number in the above range!')
+    sys.exit(1)
 
 # Create TUN device for network capture and injections
 tun = TunTapDevice(name='teletun-device')
 
-print tun.name + ' has been created, information follows:'
+print(tun.name + ' has been created, information follows:')
 
 
 # Set IP address based on --server flag
@@ -65,10 +60,10 @@ else:
 tun.netmask = '255.255.255.0'
 tun.mtu = 1500
 
-print 'Address: ' + tun.addr
-print 'Dest.-Address: ' + tun.dstaddr
-print 'Netmask: ' + tun.netmask
-print 'MTU: ' + str(tun.mtu)
+print('Address: ' + tun.addr)
+print('Dest.-Address: ' + tun.dstaddr)
+print('Netmask: ' + tun.netmask)
+print('MTU: ' + str(tun.mtu))
 
 
 # Start TUN device
@@ -83,7 +78,6 @@ received = 0
 
 # Helper function that can be executed in a thread
 def main_loop_starter():
-    pass
     receiver.start()
     # Start the receive loop
     receiver.message(main_loop())
@@ -93,60 +87,53 @@ def main_loop_starter():
 def main_loop():
     global up
     global received
-    try:
-        while up:
-            # Receive message from telegram, this includes ALL messages
-            msg = (yield)
-            # Check if it is an actual "message" message and if the sender is our peer
-            if msg is not None and msg['event'] == unicode('message')\
-                    and not msg['own'] and msg['sender']['peer_id'] == peer_id:
-                try:
-                    # Decode data and write it to the tunnel
-                    data = base64.b64decode(msg.text)
-                    received += len(data)
-                    tun.write(data)
-                except:
-                    print 'Invalid Base64-Data: ' + msg.text
-    except:
-        print '!!! Receiver crashed!'
+    while up:
+        # Receive message from telegram, this includes ALL messages
+        msg = (yield)
+        # Check if it is an actual "message" message and if the sender is our peer
+        if (
+            msg is not None and
+            msg['event'] == unicode('message') and
+            not msg['own'] and
+            msg['sender']['peer_id'] == peer_id
+        ):
+            # Decode data and write it to the tunnel
+            data = base64.b64decode(msg.text)
+            received += len(data)
+            tun.write(data)
 
 
-print 'TUN is up'
+print('TUN is up')
 
 
 # Create the receive thread via our helper method
 thread = threading.Thread(target=main_loop_starter)
 
-try:
+# Start the thread for receiving
+print('Connecting to peer...')
+thread.start()
+print('Connected! Sending Invitation!')
 
-    # Start the thread for receiving
-    print 'Connecting to peer...'
-    thread.start()
-    print 'Connected! Sending Invitation!'
+# Send the invitation message
+sender.msg(username, unicode('Hello, I would like to establish a Layer 3 Tunnel with you! -teletun'))
 
-    # Send the invitation message
-    sender.msg(username, unicode('Hello, I would like to establish a Layer 3 Tunnel with you! -teletun'))
-
-    while up:
-        # Continually read from the tunnel and write data to telegram in base64
-        # TODO: Telegram supports unicode, base64 can probably be replaced for something less overhead-inducing
-        buf = tun.read(tun.mtu)
-        data = base64.b64encode(buf)
-        sent += len(data)
-        sender.msg(username, unicode(data))
-
-except:
-    print '(Interrupt/Crash)'
+while up:
+    # Continually read from the tunnel and write data to telegram in base64
+    # TODO: Telegram supports unicode, base64 can probably be replaced for something less overhead-inducing
+    buf = tun.read(tun.mtu)
+    data = base64.b64encode(buf)
+    sent += len(data)
+    sender.msg(username, unicode(data))
 
 # Cleanup and stop application
 up = False
 tun.down()
 receiver.stop()
 
-print 'Bytes sent via Telegram: ' + str(sent)
-print 'Bytes received via Telegram: ' + str(received)
+print('Bytes sent via Telegram: ' + str(sent))
+print('Bytes received via Telegram: ' + str(received))
 
-print '~~ Bye bye! ~~'
+print('~~ Bye bye! ~~')
 
 # Literally Overkill