#!/usr/local/bin/ruby
#                                                      2001/07/04
#
#           
#                       Takashi Shimizu <simtak@mua.biglobe.ne.jp>
#                                          (c)2001 Takashi Shimizu

# 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, 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.

require "socket"
require "getopts"

Service={"http"=>80}
def parseurl(url)  #=> [protocol, host, port, location]
  url.scan(%r!(\w+)://([^/:]+)(:\d+)?(\S*)!){|protocol,host,port,loc| #/
    loc="/" unless loc
    if port;port=port[1..-1].to_i;else port=Service[protocol];end
    return protocol,host,port,loc
  }
end

class HTTPGet
  def initialize(sock)
    @sock=sock
    @header={}
    @commandline=""
    @headerdata=""
    @body=""
    @command="";   @url=""; @httpversion=""; @protocol=""
    @host=""; @port=""; @location=""
    @buffer="" ;@prompt; @peer=nil
  end    
  attr_accessor :host, :port, :prompt, :peer, :url, :header

  def HTTPGet.readsock(sock, prompt=">", peer=nil)
    http=new(sock)
    http.peer=peer
    http.prompt=prompt
    http.getheader{|line| 
      yield line if block_given? 
    }
    http.getbody{|line| 
      yield line if block_given? 
    }
    http
  end
  def getheader
    while not (@commandline=@sock.gets); ;end
    STDERR.print(@prompt, @commandline) unless $quiet
    yield @commandline if block_given?
    while line=@sock.gets
      @headerdata<<line
      yield line if block_given?
      STDERR.print(@prompt, line) unless $quiet
      break if line=~/^\r?\n/
    end
    header_analysis()
  end

  def getbody
    contentlength=@header["content-length"]
    unless @peer
      @body=@sock.read(contentlength.to_i)
    else      
      if @header["content-type"]=~/text/
	@body=readlength_with_line(contentlength){|line|
	  yield line if block_given?
	}
      else
	@body=readlength(contentlength){|line|
	  yield line if block_given?
	}
      end
    end
  end
  def to_remote
    @commandline=@command+" "+@location+" "+@httpversion+"\r\n"
    self
  end

  def to_s
    @commandline+@headerdata+@body
  end

  private
  def header_analysis
    @commandline.scan(/^(GET|POST|HEAD)\s*(\S+)\s+(\S+)/){
      |command1,url1,httpversion1|
      @command=command1; @url=url1; @httpversion=httpversion1
      @protocol, @host, @port, @location = parseurl(@url)
    }
    @headerdata.split("\n").each do |line|
      d=line.split("\s*:\s*")
      @header[d[0].downcase]=d[1].to_s.sub(/\r/,"")
    end
  end
  def readlength(length, unit=1024)
    buf="";buf0=""
    sizenow=0
    if length
      length=length.to_i
      while sizenow+unit < length
	buf0=@sock.read(unit)
	if buf0
	  yield buf0 if block_given?
	  buf<<buf0
	  sizenow+=buf0.size
	end
      end
      buf0=@sock.read(length-sizenow)
      if buf0
	yield buf0 if block_given?
	buf<<buf0
      end
    else
      while buf0=@sock.read(unit)
	yield buf0 if block_given?
	buf<<buf0
      end
    end
    buf    
  end
  def readlength_with_line(length)
    buf="";line=""
    length=length.to_i if length
    @sock.each_line{|line|
      buf<<line
      yield line if block_given?
      break if length && (length-=line.size)<=0
    }
    buf
  end
end

class Proxy
  def initialize(port)
    @extraproxy=nil
    @extraproxyport=nil
    @server=TCPServer.open(port)
  end
  attr_accessor :extraproxy, :extraproxyport

  def Proxy.usageexit()
    print "proxy.rb [--proxy proxyserver] port\n"
    print " -q : quiet\n"
    print " --proxy proxyurl : proxy via proxy\n"
    print " --daemon : run as daemon\n"
    print " ex. proxy.rb --proxy http://localhost:8080/ 8082\n"
    exit 0
  end

  def Proxy.main
    if getopts("dqh","proxy:","daemon")==nil
      raise "unknown option"
    end
    if $OPT_h;usageexit();end
    $DEBUG=$OPT_d
    $quiet=$OPT_q
    Thread.abort_on_exception=true # if $DEBUG
    port=ARGV.shift
    begin
      proxy=new(port)
      daemon if $OPT_daemon
      if $OPT_proxy
	pt,proxy.extraproxy, proxy.extraproxyport=parseurl($OPT_proxy.to_s)
      end
      proxy.mainloop
    rescue Errno::EADDRINUSE
      $stderr.print("port #{port} is already used.\n")
      exit(1)
    rescue
      $stderr.printf("%s: %s\n\t%s\n", $0, $!, $@.join("\n\t"))
      exit(1)
    end
  end
  def Proxy.daemon
    $quiet=true
    exit! if fork
    Process::setsid
    exit! if fork
    Dir::chdir("/")
    File::umask(022)
    STDIN.reopen(open("/dev/null", "r"))
    STDOUT.reopen(open("/dev/null", "w"))
    STDERR.reopen(open("/dev/null", "w"))
  end

  def mainloop
    thparent=Thread.current
    while true
      Thread.start(@server.accept) do |sock|
	th=Thread.current.id
	STDERR.print "#{th} start\n" unless $quiet
	begin
	  req=HTTPGet.readsock(sock,"#{th} < ",nil)
	  if @extraproxy
	    sockr=TCPSocket.open(@extraproxy, @extraproxyport)
	    sockr.print req.to_s
	  else
	    sockr=TCPSocket.open(req.host,req.port)
	    sockr.print req.to_remote.to_s
	  end
	  HTTPGet.readsock(sockr,"#{th} > ",req.url){|str|
	      sock.print str
	  }
	rescue SystemCallError
	  STDERR.print "#{th} terminated with #{$!}\n" unless $quiet
	rescue
	  STDERR.puts $! unless $quiet
	ensure
	  sockr.close if sockr	
	  sock.close if sock
	  STDERR.print "#{th} finished\n" unless $quiet
	  if (not $quiet) && (stack=ThreadGroup::Default.list).size>2
	    stack=(stack-[thparent,Thread.current]).map{|t|t.id}
	    STDERR.print("["+stack.join(" ")+"]\n")
	  end
	  STDERR.flush unless $quiet
	end
      end
    end
  end
end

Proxy.main()