Code cleanup

.gitignore

@@ -87,3 +87,6 @@ ENV/
 
 # Rope project settings
 .ropeproject
+
+# JetBrains
+.idea

teletun.py

@@ -1,31 +1,37 @@
 #!/usr/bin/env python
 
+# (C) Stefan Reiter 2017
+
 from pytun import TunTapDevice
 from pytg.sender import Sender
 from pytg.receiver import Receiver
 from pytg.utils import coroutine
+from pytg.exceptions import ConnectionError
 import base64
 import threading
 import sys
-import traceback
-import os
-import signal
-import psutil
 
 
+# Connect to telegram
 print 'Connecting to telegram...'
 
 receiver = Receiver(host="localhost", port=4458)
 sender = Sender(host="localhost", port=4458)
 
+# Retrieve contact list
 contacts = {}
 
-counter = 0
+try:
-for c in sender.dialog_list():
+    counter = 0
+    for c in sender.dialog_list():
         counter += 1
         contacts[counter] = c
         print unicode(counter) + ': \t' + unicode(c['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)
 
+# Ask user to choose contact
 i = input('Telegram online, please enter contact to connect to (by number): ')
 
 if i not in contacts:
@@ -33,15 +39,18 @@ if i not in contacts:
     exit(1)
 
 
+# Print username
 username = unicode(contacts[i]['print_name'])
 print 'Connecting to partner: ' + username
 
 
+# Create TUN device for network capture and injections
 tun = TunTapDevice(name='teletun-device')
 
 print tun.name + ' has been created, information follows:'
 
 
+# Set IP address based on --server flag
 if '--server' in sys.argv:
     tun.addr = '10.8.0.1'
     tun.dstaddr = '10.8.0.2'
@@ -57,13 +66,17 @@ print 'Dest.-Address: ' + tun.dstaddr
 print 'Netmask: ' + tun.netmask
 print 'MTU: ' + str(tun.mtu)
 
+
+# Start TUN device
 tun.up()
 up = True
 
 
+# Helper function that can be executed in a thread
 def main_loop_starter():
     pass
     receiver.start()
+    # Start the receive loop
     receiver.message(main_loop())
 
 
@@ -72,60 +85,49 @@ def main_loop():
     global up
     try:
         while up:
+            # Receive message from telegram, this includes ALL messages
             msg = (yield)
-            if msg is not None and msg['event'] == unicode('message') and not msg['own']:
+            # 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['from']['print_name'] == username:
                 try:
+                    # Decode data and write it to the tunnel
                     data = base64.b64decode(msg.text)
                     tun.write(data)
                 except:
                     print msg.text
     except:
-        exc_type, exc_value, exc_traceback = sys.exc_info()
-        lines = traceback.format_exception(exc_type, exc_value, exc_traceback)
-        print ''.join('!! ' + line for line in lines)
         print 'Receiver stopped'
 
 
 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!'
 
+    # Send the invitation message
     sender.msg(username, unicode('Hello, I would like to establish a Layer 3 Tunnel with you! -teletun'))
 
-    while True:
+    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)
         sender.msg(username, unicode(base64.b64encode(buf)))
 
 except:
-    exc_type, exc_value, exc_traceback = sys.exc_info()
-    lines = traceback.format_exception(exc_type, exc_value, exc_traceback)
-    print ''.join('!! ' + line for line in lines)
     print 'Exiting...'
 
+# Cleanup and stop application
 up = False
 tun.down()
-
 receiver.stop()
 
 print 'Bye bye!'
-
-
-# Literally Overkill
-
-current_process = psutil.Process()
-children = current_process.children(recursive=True)
-for child in children:
-    os.kill(child.pid, signal.SIGKILL)
-
-children = current_process.children(recursive=True)
-for child in children:
-    os.kill(child.pid, signal.SIGKILL)
-
-os.kill(current_process.pid, signal.SIGKILL)