(* Interface to the Hostip.info service (see http://hostip.info).
 * Copyright (C) 2005-2006 Merjis Ltd.
 * $Id: hostIP.ml,v 1.2 2005/12/14 15:02:59 rich Exp $
 *)

open Printf
open ExtString

type t = {
  connection : Http_client.pipeline;
  cache : (string, result) Hashtbl.t;
  cache_file : string option;
}

and result = {
  country_code : string option;
  country_name : string option;
  city : string option;
  location : (string * string) option;
}

let is_unknown s =
  s = "" ||
  String.starts_with s "(Unknown" ||
  String.starts_with s "(Private"

let perm = 0o666

let load_cache_file cache filename =
  try
    (* Assume the file doesn't exist and try to create it. *)
    let flags = [ Unix.O_WRONLY; Unix.O_CREAT; Unix.O_EXCL ] in
    let fd = Unix.openfile filename flags perm in
    Unix.close fd
  with
    Unix.Unix_error _ ->
      (* File exists, open for reading, acquire a lock, and read the
       * contents into the cache.
       *)
      let chan = open_in filename in
      let fd = Unix.descr_of_in_channel chan in
      Unix.lockf fd Unix.F_RLOCK 0;
      let rec loop () =
	  let line = input_line chan in
	  (* Split into elements. *)
	  let line = String.nsplit line "|" in
	  (* First element is the version number.  Ignore higher versions
	   * than we currently understand.
	   *)
	  (match line with
	   | [ "1";
	       ip; country_code; country_name; city; latitude; longitude ] ->
	       let country_code =
		 if country_code = "" then None else Some country_code in
	       let country_name =
		 if country_name = "" then None else Some country_name in
	       let city = if city = "" then None else Some city in
	       let latitude = if latitude = "" then None else Some latitude in
	       let longitude =
		 if longitude = "" then None else Some longitude in
	       let location =
		 match latitude, longitude with
		 | Some latitude, Some longitude -> Some (latitude, longitude)
		 | _ -> None in
	       let result = {
		 country_code = country_code; country_name = country_name;
		 city = city; location = location } in
	       Hashtbl.replace cache ip result
	   | _ -> () (* ignore *)
	  );
	  loop () in
      (try loop () with End_of_file -> ());
      close_in chan

let append_cache_file ip result filename =
  (* Open the file and lock for writing, non-destructively. *)
  let flags = [ Open_wronly; Open_creat ] in
  let chan = open_out_gen flags perm filename in
  let fd = Unix.descr_of_out_channel chan in
  Unix.lockf fd Unix.F_LOCK 0;

  (* Seek to the end of the file. *)
  seek_out chan (out_channel_length chan);

  (* Append result. *)
  fprintf chan "1|%s|%s|%s|%s|%s|%s\n"
    ip
    (match result.country_code with None -> "" | Some s -> s)
    (match result.country_name with None -> "" | Some s -> s)
    (match result.city with None -> "" | Some s -> s)
    (match result.location with None -> "" | Some (s, _) -> s)
    (match result.location with None -> "" | Some (_, s) -> s);

  close_out chan

let connection ?connection ?cache_file () =
  let connection =
    match connection with
    | None ->
	let event_system = Unixqueue.create_unix_event_system () in
	let connection = new Http_client.pipeline in
	connection#set_event_system event_system;
	connection
    | Some conn -> conn in
  let cache = Hashtbl.create 127 in
  (match cache_file with
   | Some filename -> load_cache_file cache filename
   | None -> ()
  );
  {
    connection = connection;
    cache = cache;
    cache_file = cache_file;
  }

let http_url_syntax = Hashtbl.find Neturl.common_url_syntax "http"

let split_lines_re = Pcre.regexp "\\r?\\n"
let country_re = Pcre.regexp "(.*)\\s+\\(([A-Z]{2})\\)$"

let get {connection = connection; cache = cache; cache_file = cache_file} ip =
  try
    Hashtbl.find cache ip
  with
    Not_found ->
      let result =
	let query =
	  Netencoding.Url.mk_url_encoded_parameters
	    ["ip", ip; "position", "true"] in
	let uri = Neturl.make_url
	  ~encoded:true
	  ~scheme:"http"
	  ~host:"api.hostip.info" ~path:[""; "get_html.php"] ~query
	  http_url_syntax in
	let uri = Neturl.string_of_url uri in
	let rq = new Http_client.get uri in
	connection#add rq;
	connection#run ();
	let res = rq#get_resp_body () in
	let lines = Pcre.split ~rex:split_lines_re res in
	let lines = List.map (fun s -> String.split s ": ") lines in
	let country =
	  try Some (List.assoc "Country" lines) with Not_found -> None in
	let country =
	  match country with
	  | None -> None
	  | Some country when is_unknown country -> None
	  | country -> country in
	let country_name, country_code =
	  match country with
	  | None -> None, None
	  | Some country ->
	      try
		let subs = Pcre.exec ~rex:country_re country in
		let sub1 = Pcre.get_substring subs 1 in
		let sub2 = Pcre.get_substring subs 2 in
		Some sub1, Some sub2
	      with
		Not_found -> None, None in

	let city =
	  try Some (List.assoc "City" lines) with Not_found -> None in
	let city =
	  match city with
	  | None -> None
	  | Some city when is_unknown city -> None
	  | city -> city in

	let latitude =
	  try Some (List.assoc "Latitude" lines) with Not_found -> None in
	let latitude =
	  match latitude with
	  | None -> None
	  | Some latitude when is_unknown latitude -> None
	  | latitude -> latitude in

	let longitude =
	  try Some (List.assoc "Longitude" lines) with Not_found -> None in
	let longitude =
	  match longitude with
	  | None -> None
	  | Some longitude when is_unknown longitude -> None
	  | longitude -> longitude in

	let location =
	  match latitude, longitude with
	  | Some latitude, Some longitude -> Some (latitude, longitude)
	  | _ -> None in

	{ country_code = country_code;
	  country_name = country_name;
	  city = city;
	  location = location } in
      Hashtbl.add cache ip result;
      (match cache_file with
       | Some filename -> append_cache_file ip result filename
       | None -> ()
      );
      result

let get_country_code conn ip =
  let result = get conn ip in
  result.country_code

let get_country_name conn ip =
  let result = get conn ip in
  result.country_name

let get_city conn ip =
  let result = get conn ip in
  result.city

let get_location conn ip =
  let result = get conn ip in
  result.location

let clear_cache { cache = cache } =
  Hashtbl.clear cache

let forget { cache = cache } ip =
  Hashtbl.remove cache ip