/PuttyKnife/dist/connect.py

#!/usr/bin/env python    
# -*- coding: utf-8 -*-
# vi: ts=4 sw=4

#-------------------------------------------------------------------------------
# Program: PuttyKnife VPN - A PuTTY/ssh-based virtual private networking client
# Module:  connection supervisor
# Desc:    Stylus Toolbox is a GTK+ GUI front-end for Gutenprint's 
#          escputil printer utility.
# Author:  Rob A. Shinn
#
# Legal:    Copyright 2008 Rob A. Shinn. All rights reserved.
#           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.
#-------------------------------------------------------------------------------

import sys
import os
import unittest
#from pexpect import *
from configuration import *
from puttysession import *

class Tunnel():
    def __init__(self, 
                listen_addr=None, 
                listen_port=None, 
                host_addr=None, 
                host_port=None, 
                remote=False):
        self.listen_addr=listen_addr
        self.listen_port=listen_port
        self.host_addr=host_addr
        self.host_port=host_port
        self.remote=remote
    def __repr__(self):
        return "%s:%s -> %s:%s (%s)" % (self.listen_addr, self.listen_port, 
                                        self.host_addr, self.host_port, 
                                        self.remote)
        
class Tunnels():
    def __init__(self):
        self.forwarders=[]
    def add_tunnel(self, tunnel):
        self.forwarders.append(tunnel)
    def get_list(self):
        return self.forwarders
    def __iter__(self):
        self.iter=self.forwarders.__iter__()
        return self
    def next(self):
        row = self.iter.next()
        return row
    
                
class Connection():
    def __init__(self, 
                 host=None, 
                 port=None, 
                 username=None,
                 password=None,
                 keyfile=None,
                 tunnels=None,
                 ip_version=None, 
                 compression=False,
                 verbose=False,
                 x11=False,
                 agent=False,
                 useagent=True,
                 shell=True,
                 socks_port=None,
                 ssh_version=None,
                 pty=True,
                 logfile=None,
                 proxy_dns=1,
                 proxy_host=None,
                 proxy_port=None,
                 proxy_method=None,
                 proxy_username=None,
                 proxy_password=None,
                 proxy_command=None,
                 ):
        self.command_line=None
        self.host=host
        self.port=port
        self.username=username
        self.password=password
        self.keyfile=keyfile
        self.tunnels=tunnels
        self.ip_version=ip_version
        self.ssh_version=ssh_version
        self.compression=compression
        self.verbose=verbose
        self.socks_port=socks_port
        self.x11=x11
        self.agent=agent
        self.pty=pty
        self.shell=shell
        self.useagent=useagent
        self.logfile=logfile
        self.proxy_host=proxy_host
        self.proxy_method=proxy_method
        self.proxy_port=proxy_port
        self.proxy_username=proxy_username
        self.proxy_password=proxy_password
        self.proxy_command=proxy_command
        self.proxy_dns=proxy_dns
        #self.create_putty_session('PuttyKnife_VPN_Session')

    def create_putty_session(self, filename):
        psf=putty_session(filename=filename,
                          proxy_method=self.proxy_method, 
                          proxy_host=self.proxy_host, 
                          proxy_port=self.proxy_port,
                          proxy_username=self.proxy_username, 
                          proxy_password=self.proxy_password, 
                          proxy_command=self.proxy_command,
                          proxy_dns=self.proxy_dns)
        psf.optionxform=str
        psf.write()
    def write_log(self, msg):
        if self.logfile:
            self.logfile.write(msg)
    def __log_filter(self, msg):
        if self.logfile:
            self.__write_log(self, msg)
        return msg
    def open(self):
        switches='-v'
        if self.keyfile:
            switches='%s -i %s' % (switches,self.keyfile)
        if self.port:
            switches='%s -P %s' % (switches,self.port)
        if self.username:
            switches='%s -l %s' % (switches,self.username)
        if self.ip_version:
            switches='%s -%s' % (switches,self.ip_version)
        if self.compression:
            switches='%s -C' % switches
        if self.socks_port:
            switches='%s -D %s' % (switches,self.socks_port)
        if self.x11:
            switches='%s -X' % switches
        if self.agent:
            switches='%s -agent' % switches
        else:
            switches='%s -noagent' % switches
        if not self.shell:
            switches='%s -N' % switches
        if not self.pty:
            switches='%s -T' % switches
        for tunnel in self.tunnels:
            if tunnel.remote:
                switches='%s -R ' % switches
            else:
                switches='%s -L ' % switches
            if tunnel.listen_addr:
                switches='%s%s:' % (switches,tunnel.listen_addr)
            switches='%s%s:' % (switches,tunnel.listen_port)
            switches='%s%s:' % (switches,tunnel.host_addr)
            switches='%s%s' % (switches,tunnel.host_port)
        if self.host:
            command_line='plink %s %s' % (switches, self.host)
        else:
            sys.quit()
            
        child=spawn(command_line)
        self.write_log('Launching command: %s' % command_line)
        if self.password:
            child.expect('password:')
            child.send(self.password)
        else:
            child.expect('Opened channel for session')
        self.write_log(child.before)
        msg=''
        while child.isalive():
            try:
                ch=child.read(size=1)
            except TIMEOUT:
                pass
            if ch=='\n':
                self.write_log(msg)
                msg=''
            else:
                 msg=msg+ch
            
class __unittest__(unittest.TestCase):
    def setUp(self):
        connection=Connection()
        connection.host='brighid'
        connection.keyfile='/home/morgan/.ssh/id_rsa.ppk'
        connection.agent=False
        connection.logfile=logfile
        tunnel=Tunnel(listen_addr='127.0.0.1', listen_port=2222, 
                      host_addr='192.168.1.101', host_port=22, remote=False)
        tunnels=Tunnels()
        tunnels.add_tunnel(tunnel)
        tunnel=Tunnel(listen_addr='127.0.0.1', listen_port=9010, 
                      host_addr='192.168.1.101', host_port=901, remote=False)
        tunnels.add_tunnel(tunnel)
        connection.tunnels=tunnels
        self.connection=connection
    def runTest(self):
        self.connection.open()
        
if __name__ == '__main__':
    from logger import *
    unittest.main()