mirror of https://github.com/nealey/irc-bot
285 lines
6.6 KiB
OCaml
285 lines
6.6 KiB
OCaml
|
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
|