irc-bot/irc.ml

285 lines
6.6 KiB
OCaml
Raw Normal View History

open Uq_engines
open Unixqueue
let newline_re = Pcre.regexp "\n\r?"
let argsep_re = Pcre.regexp " :"
let space_re = Pcre.regexp " "
let string_map f s =
let l = String.length s in
if l = 0 then
s
else
let r = String.create l in
for i = 0 to l - 1 do
String.unsafe_set r i (f (String.unsafe_get s i))
done;
r
let lowercase_char c =
if (c >= 'A' && c <= '^') then
Char.unsafe_chr(Char.code c + 32)
else
c
let uppercase_char c =
if (c >= 'a' && c <= '~') then
Char.unsafe_chr(Char.code c - 32)
else
c
let uppercase s = string_map uppercase_char s
let lowercase s = string_map lowercase_char s
class irc (ues : unix_event_system) =
object (self)
val mutable debug = false
val throttle_interval = 1.0
(** Group for this bot's events *)
val g = ues#new_group ()
val out_wait_id = ues#new_wait_id ()
val mutable check_output = None
val mutable output_pending = ""
val mutable input_pending = ""
val outq_immediate = Queue.create ()
val outq_throttled = Queue.create ()
val outq_last_sent = Unix.time ()
initializer
ues#add_handler g self#handle_event
method debug v =
debug <- v
method log msg =
if debug then
prerr_endline msg
method handle_event ues' esys e =
assert (ues' = ues);
match e with
| Input_arrived (g, fd) ->
self#handle_input fd
| Output_readiness (g, fd) ->
self#handle_output fd
| Out_of_band (g, fd) ->
self#handle_oob fd
| Timeout (g, op) ->
self#handle_timeout op
| Signal ->
self#handle_signal ()
| Extra exn ->
self#handle_extra exn
method handle_input fd =
let s = 4096 in
let buf = String.create s in
let len = Unix.read fd buf 0 s in
let input = input_pending ^ (String.sub buf 0 len) in
if (input <> "") then
let lines = Pcre.split ~rex:newline_re input in
let rec handle_lines lines =
match lines with
| []
| [""] ->
input_pending <- "";
| line :: tl ->
self#handle_line line;
handle_lines tl
in
handle_lines lines
else
begin
Unix.close fd;
ues#clear g;
end
method handle_output fd =
let data = (
if (output_pending <> "") then
output_pending
else if not (Queue.is_empty outq_immediate) then
Queue.pop outq_immediate
else if not (Queue.is_empty outq_throttled) then
begin
(* Stop listening for output events; add a timeout to
start listening again *)
(match check_output with
| None -> ()
| Some co ->
let pay_attention () =
ues#add_resource g (co, -.1.0)
in
ues#remove_resource g co;
ues#once g throttle_interval pay_attention
);
Queue.pop outq_throttled
end
else
match check_output with
| None -> ""
| Some co ->
ues#remove_resource g co;
""
) in
let data_len = String.length data in
let n = Unix.single_write fd data 0 data_len in
output_pending <- String.sub data n (data_len - n);
if (data <> "") then
self#log ("--> " ^ (String.escaped data))
method handle_oob fd =
self#log "OOB ready!";
raise Equeue.Reject
method handle_timeout op =
self#log "Timeout!";
raise Equeue.Reject
method handle_signal () =
self#log "Signal!";
raise Equeue.Reject
method handle_extra exn =
self#log "Extra!";
raise Equeue.Reject
method handle_line line =
let argstr, txt =
match Pcre.split ~max:2 ~rex:argsep_re line with
| [] -> ("", "")
| [a] -> (a, "")
| [a; b] -> (a, b)
| _ -> ("", "")
in
let sender, args =
let args' = Pcre.split ~rex:space_re argstr in
if (debug) then
print_endline ("<-- [" ^
(String.concat "; " args') ^
"] " ^
txt);
if (List.hd args').[0] = ':' then
(List.hd args', List.tl args')
else
("", args')
in
match args with
| [] ->
()
| "NOTICE" :: args ->
let tlen = String.length txt in
if ((txt.[0] = '\001') &&
(txt.[tlen - 1] = '\001')) then
self#handle_ctcp_reply sender args (String.sub txt 1 (tlen - 2))
else
self#handle_notice sender args txt
| "PRIVMSG" :: args ->
let tlen = String.length txt in
if ((txt.[0] = '\001') &&
(txt.[tlen - 1] = '\001')) then
self#handle_ctcp_request sender args (String.sub txt 1 (tlen - 2))
else
self#handle_privmsg args sender txt
| ["PING"] ->
self#handle_ping txt
| str :: args ->
let numeric =
try
Some (int_of_string str)
with Failure _ ->
None
in
match numeric with
| Some n ->
self#handle_numeric n sender args txt
| None ->
self#handle_unknown sender args txt
method handle_privmsg sender args txt =
()
method handle_notice sender args txt =
()
method handle_ctcp_request sender args txt =
()
method handle_ctcp_reply sender args txt =
()
method handle_numeric n sender args txt =
()
method handle_unknown sender args txt =
self#log ("Got unknown server message")
method handle_ping txt =
self#send ["PONG"] txt
(** Public methods *)
method set_fd fd nick gecos =
(* XXX: Clear old junk *)
check_output <- Some (Wait_out fd);
ues#add_resource g (Wait_in fd, -.1.0);
self#send ~now:true ["USER"; nick; "+iw"; nick] gecos;
self#send ~now:true ["NICK"; nick] ""
(** Send a command to the IRC server *)
method send ?(now=false) args txt =
match check_output with
| None -> ()
| Some co ->
let q = (if now then outq_immediate else outq_throttled) in
let cmdstr =
(String.concat " " args) ^
(if txt = "" then "" else " :") ^
txt ^
"\n"
in
Queue.push cmdstr q;
ues#add_resource g (co, -.1.0)
(** Send a private message *)
method privmsg ?(now=false) recipient txt =
self#send ~now ["PRIVMSG"; recipient] txt
(** Send a notice *)
method notice ?(now=false) recipient txt =
self#send ~now ["NOTICE"; recipient] txt
(** Send a CTCP request *)
method ctcp_request ?(now=false) recipient command txt =
self#privmsg ~now recipient ("\001" ^ command ^ " " ^ txt ^ "\001")
(** Send a CTCP reply *)
method ctcp_reply ?(now=false) recipient command txt =
self#notice ~now recipient ("\001" ^ command ^ " " ^ txt ^ "\001")
end
let main() =
let ues = new unix_event_system () in
let c = connector (`Socket(`Sock_inet_byname(Unix.SOCK_STREAM,
"woozle.org", 6667),
default_connect_options
)) ues in
when_state
~is_done:(fun connstat ->
match connstat with
| `Socket(fd, _) ->
let b = new irc ues in
b#set_fd fd "plasma" "Plasma Bot"
| _ -> assert false
)
c;
Unixqueue.run ues