gmail_transport.py :  » Email » libgmail » libgmail-0.1.11 » Python Open Source

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# ----------------------------------------------------------------------------------
# Copyleft (K) by Jose Rodriguez. This source is free (GPL)
# Partially based on John Nielsen ASPN recipe (http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/301740)
# Partially based on Alessandro Budai recipe (http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/456195)
# ----------------------------------------------------------------------------------


# ClientCookie to connection through a proxy using the CONNECT method, (useful for SSL)
# tested with python 2.4

import mechanize as ClientCookie
import urllib
import httplib
import socket
import base64


def split_proxy_URL(proxy):
  if proxy is None:
      return None, None, None

  try:
      if proxy[:7] != 'http://':  # Ensures proxy string begins with 'http://'
          proxy = 'http://' + proxy
  except:
      pass

  proxy_username = proxy_password = None

  urltype, r_type = urllib.splittype(proxy)
  proxy, XXX = urllib.splithost(r_type)
  if '@' in proxy:
      proxy_username, proxy = proxy.split('@', 1)
      if ':' in proxy_username:
          proxy_username, proxy_password = proxy_username.split(':', 1)

  return proxy, proxy_username, proxy_password



class ProxyHTTPConnection(httplib.HTTPConnection):

  _ports = {'http' : 80, 'https' : 443}

  def request(self, method, url, body=None, headers={}):
    #request is called before connect, so can interpret url and get
    #real host/port to be used to make CONNECT request to proxy
    proto, rest = urllib.splittype(url)
    if proto is None:
      raise ValueError, "unknown URL type: %s" % url

    host, rest = urllib.splithost(rest) # get host
    host, port = urllib.splitport(host) #try to get port

    #if port is not defined try to get from proto
    if port is None:
      try:
        port = self._ports[proto]
      except KeyError:
        raise ValueError, "unknown protocol for: %s" % url

    self._real_host = host
    self._real_port = port
    httplib.HTTPConnection.request(self, method, url, body, headers)
    

  def connect(self):
    httplib.HTTPConnection.connect(self)

    self.send("CONNECT %s:%d HTTP/1.0\r\n" % (self._real_host, self._real_port))
    if self.proxy_user is not None and self.proxy_passwd is not None:
      cred = base64.encodestring("%s:%s" % (urllib.unquote(self.proxy_user), urllib.unquote(self.proxy_passwd))).strip()
      self.send("Proxy-authorization: Basic %s\r\n" % cred)

    self.send("User-Agent: Mozilla/5.0 (Compatible; libgmail-python)\r\n\r\n")
    response = self.response_class(self.sock, strict=self.strict, method=self._method)
    (version, code, message) = response._read_status()
    #probably here we can handle auth requests...
    if code != 200:
      #proxy returned and error, abort connection, and raise exception
      self.close()
      raise socket.error, "Proxy connection failed: %d %s" % (code, message.strip())

    #eat up header block from proxy....
    while True:
      line = response.fp.readline() #should not use directly fp probablu
      if line == '\r\n': break


  @classmethod
  def new_auth(cls, proxy_host, proxy_user = None, proxy_passwd = None):
    cls.proxy_host = proxy_host
    cls.proxy_user = proxy_user
    cls.proxy_passwd = proxy_passwd

    return cls



class ProxyHTTPSConnection(ProxyHTTPConnection):
  
  default_port = 443

  def __init__(self, host, port = None, key_file = None, cert_file = None, strict = None):
    ProxyHTTPConnection.__init__(self, host, port)
    self.key_file = key_file
    self.cert_file = cert_file
  
  def connect(self):
    ProxyHTTPConnection.connect(self)
    #make the sock ssl-aware
    ssl = socket.ssl(self.sock, self.key_file, self.cert_file)
    self.sock = httplib.FakeSocket(self.sock, ssl)

    
class ConnectHTTPHandler(ClientCookie.HTTPHandler):

  def __init__(self, proxy=None, debuglevel=0):
    self.proxy, self.proxy_user, self.proxy_passwd = split_proxy_URL(proxy)
    ClientCookie.HTTPHandler.__init__(self, debuglevel)

  def do_open(self, http_class, req):
    if self.proxy is not None:
      req.set_proxy(self.proxy, 'http')
    return ClientCookie.HTTPHandler.do_open(self, ProxyHTTPConnection.new_auth(self.proxy, self.proxy_user, self.proxy_passwd), req)
  


class ConnectHTTPSHandler(ClientCookie.HTTPSHandler):

  def __init__(self, proxy=None, debuglevel=0):
    self.proxy, self.proxy_user, self.proxy_passwd = split_proxy_URL(proxy)
    ClientCookie.HTTPSHandler.__init__(self, debuglevel)

  def do_open(self, http_class, req):
    if self.proxy is not None:
      req.set_proxy(self.proxy, 'https')
    return ClientCookie.HTTPSHandler.do_open(self, ProxyHTTPSConnection.new_auth(self.proxy, self.proxy_user, self.proxy_passwd), req)