(* DCC module.

This file is part of Liece.                                          

Author: Daiki Ueno <daiki@kake.info.waseda.ac.jp>                    
Created: 1998-09-28                                               
Revised: 1999-01-28                                               
Keywords: IRC, liece, DCC                                        

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.                         
                                                                      
You should have received a copy of the GNU General Public License    
along with GNU Emacs; see the file COPYING.  If not, write to the    
Free Software Foundation, Inc., 59 Temple Place - Suite 330,         
Boston, MA 02111-1307, USA.  *)

open Unix

let usage prefix progname = 
  String.concat "\n"
    (List.map (fun x -> (Printf.sprintf "%s: %s %s" prefix progname x))
       [ "send <port> <filename>";
	 "receive <host> <port> <size> <filename>";
         "chat listen <port>"; 
         "chat connect <host> <port>" ])

let buff = String.create 1024

let print_exc exc =
  match exc with
    Unix_error (err, fun_name, arg) ->
      prerr_string "\"";
      prerr_string fun_name;
      prerr_string "\" failed";
      if String.length arg > 0 then
        begin
          prerr_string " on \""; prerr_string arg; prerr_string "\""; ()
        end;
      prerr_string ": ";
      prerr_endline (error_message err);
      flush Pervasives.stderr; ()
  | _ ->
      try Printexc.print raise exc with
        _ -> ()

let accept_connection f s =
  let (t, addr) = accept s in
  f t; close t; ()
	
let write_file filename size t =
  let fd =
    try openfile filename [ O_RDONLY ] 0 with
      _ -> Printf.eprintf "Open failed.\n"; flush Pervasives.stderr; exit 1
  in
  let (len, rlen) = ref 0, ref 0 in
  while len := read fd buff 0 (String.length buff); !len <> 0
  do
    let rec loop i =
      let j = i + write t buff i (!len - i) in 
      if j == !len then () else loop j
    in
    loop 0;
    flush (out_channel_of_descr t);
    rlen := !rlen + !len;
    Printf.printf "DCC %s %d%% (%d/%d bytes) sent.\n"
      filename (100 * !rlen / size) !rlen size;
    flush Pervasives.stdout
  done;
  close fd; close t; ()

let send_file port filename =
  try
    let host = gethostbyname (gethostname ()) in
    let haddr = string_of_inet_addr host.h_addr_list.(0) in
    let s = socket PF_INET SOCK_STREAM 0 in
    setsockopt s SO_REUSEADDR true;
    bind s (ADDR_INET (inet_addr_any, port));
    let port =
      match (getsockname s) with
	ADDR_INET (addr, port) -> port
      | _ -> port
    in
    listen s 1;
    let fd =
      try openfile filename [ O_RDONLY ] 0 with
      	_ -> Printf.eprintf "Open failed.\n"; flush Pervasives.stderr; exit 1
    in
    let size = (fstat fd).st_size in
    close fd;
    Printf.printf "DCC send %s %d %s %d\n"
      (Filename.basename filename) port (Naddr.encode haddr) size;
    flush Pervasives.stdout;
    accept_connection (fun t -> write_file filename size t) s;
  with
    exc -> print_exc exc

let read_file filename size t =
  let fd =
    try openfile filename [ O_WRONLY; O_CREAT ] 0o600 with
      _ -> Printf.eprintf "Open failed.\n"; flush Pervasives.stderr; exit 1
  in
  let (len, rlen) = ref 0, ref 0 in
  while len := read t buff 0 (String.length buff); !len <> 0
  do
    let _ = write fd buff 0 !len in
    flush (out_channel_of_descr fd);
    rlen := !rlen + !len;
    Printf.printf "DCC %s %d%% (%d/%d bytes) received.\n"
      filename (100 * !rlen / size) !rlen size;
    flush Pervasives.stdout
  done;
  close fd; close t

let receive_file host port size filename =
  let s = socket PF_INET SOCK_STREAM 0 in
  connect s (ADDR_INET (inet_addr_of_string (Naddr.decode host), port));
  read_file filename size s; ()

let chat_loop s =
  let sel = ref [s; stdin] in
  while !sel <> [] do
    let (l, _, _) =
    try
      select !sel [] [] (-1.0) 
    with 
      _ -> exit 0
    in
    List.iter
      (fun x -> 
	let (rfd, wfd) =
	  if x == s then 
	    (s, stdout)
	  else if x == stdin then 
	    (stdin, s)
	  else 
	    (stdin, stdout)
	in
	let len = ref 0 
	in
	len := read rfd buff 0 (String.length buff); 
	if !len == 0 then 
	  begin try shutdown wfd SHUTDOWN_SEND; () with
	    _ -> exit 0
	  end
	else
	  begin
	    let rec loop i =
	      let j = i + write wfd buff i (!len - i) in 
	      if j == !len then () else loop j
      	    in
      	    loop 0;
	    flush (out_channel_of_descr wfd)
	  end;
	())
      l
  done

let chat_listen port =
  let host = gethostbyname (gethostname ()) in
  let haddr = string_of_inet_addr host.h_addr_list.(0) in
  let s = socket PF_INET SOCK_STREAM 0 in
  setsockopt s SO_REUSEADDR true;
  bind s (ADDR_INET (inet_addr_any, port));
  let port =
    match (getsockname s) with
      ADDR_INET (addr, port) -> port
    | _ -> port
  in
  listen s 1;
  Printf.printf "DCC chat %s %d\n" (Naddr.encode haddr) port;
  flush Pervasives.stdout;
  accept_connection 
    (fun t -> 
      Printf.printf "DCC chat established\n";
      flush Pervasives.stdout; 
      chat_loop t) s; ()

let chat_connect host port =
  let s = socket PF_INET SOCK_STREAM 0 in
  connect s (ADDR_INET (inet_addr_of_string (Naddr.decode host), port));
  Printf.printf "DCC chat established\n";
  flush Pervasives.stdout;
  chat_loop s; ()
  
let getaddr_ext server =
  let addr =
    try (gethostbyname server).h_addr_list.(0) with
      _ -> inet_addr_of_string "198.41.0.4"
  and port = 7 in 
  let s = socket PF_INET SOCK_DGRAM 0 in
  connect s (ADDR_INET (addr, port));
  match (getsockname s) with
    ADDR_INET (addr, port) -> addr
  | _ -> raise Not_found
  
let main () =
  let a = ref [] in
  let usage = usage "Usage" (Filename.basename Sys.argv.(0)) in
  let speclist = [] in
  Arg.parse speclist (fun x -> a := !a @ [x]) usage;
  begin match !a with
    "send" :: [ port; filename ] -> 
      let port = 
	try int_of_string port with
	  _ -> Arg.usage speclist usage; exit 1
      in
      send_file port filename;
      exit 0; ()
  | "receive" :: [ host; port; size; filename ] ->
      let (port, size) = 
	try 
	  int_of_string port, 
	  int_of_string size
	with
	  _ -> Arg.usage speclist usage; exit 1
      in
      receive_file host port size filename;
      exit 0; ()
  | "chat" :: [ "listen"; port ] ->
      let port =
	try
	  int_of_string port
	with
	  _ -> Arg.usage speclist usage; exit 1
      in
      chat_listen port;
      exit 0; ()
  | "chat" :: [ "connect"; host; port ] ->
      let port =
	try
	  int_of_string port
	with
	  _ -> Arg.usage speclist usage; exit 1
      in
      chat_connect host port;
      exit 0; ()
  | _ -> Arg.usage speclist usage; exit 1
  end

let _ = Printexc.catch main ()