From 33974047be172cbd38f095d48cf27fc8aa8bfde7 Mon Sep 17 00:00:00 2001 From: Epic Morphism Date: Sun, 19 May 2019 11:38:55 +0200 Subject: [PATCH] rewrite the connections functions and restructure the code LOL --- Makefile | 4 +- connection.ml | 96 ++++++++++++++++++++++++++++++++++++++++++ connection.mli | 38 +++++++++++++++++ vantabot.ml | 110 +++++++++++++++++-------------------------------- 4 files changed, 173 insertions(+), 75 deletions(-) create mode 100644 connection.ml create mode 100644 connection.mli diff --git a/Makefile b/Makefile index 2b590d0..0fc6546 100644 --- a/Makefile +++ b/Makefile @@ -1,9 +1,9 @@ -SOURCES = irccolors.ml store.ml vantabot.ml +SOURCES = irccolors.ml store.ml connection.ml vantabot.ml RESULT = vantabot PACKS = lwt,irc-client,irc-client-lwt,irc-client-tls,dokeysto THREADS = yes UNIX = yes -all: byte-code native-code +all: native-code -include OCamlMakefile diff --git a/connection.ml b/connection.ml new file mode 100644 index 0000000..95090b8 --- /dev/null +++ b/connection.ml @@ -0,0 +1,96 @@ +open Irc_client +open Lwt + +let formatf = Printf.sprintf + +(** We distinguish between "big channels" (allowed to spam) and "smol channels" (shouldn't spam). + Thanks to nilaky for the terminology. *) +module Conn (C: CLIENT with type 'a Io.t = 'a Lwt.t) = struct + type t = C.connection_t + + let smolchannels = ref [] + let bigchannels = ref [] + + let send_msg (connection : t) ~is_big ?channel (m : string) = + let open Lwt_list in + let f c = + C.send_privmsg ~connection ~target:c ~message:m + in + let channels = match channel with + | None -> if is_big then !bigchannels else !bigchannels @ !smolchannels + | Some c -> if is_big && not (List.mem c !bigchannels) then [] else [c] + in + iter_s f channels + + let insert_channel c is_big = + let r = (if is_big then bigchannels else smolchannels) in + if List.mem c !r then () + else r := c::!r + + let join_chan (connection : t) ~is_big ~channel = + C.send_join ~connection ~channel >>= fun () -> + return (insert_channel channel is_big) + + let callbacks = ref [] + + let nick_of_irc_message imsg = + let open Irc_message in + match imsg.prefix with + | None -> "" + | Some wh -> + match String.split_on_char '!' wh with + | a::_ -> a + | [] -> "" + let register_callback ~prefix ~callback = + callbacks := (prefix,callback)::!callbacks + + let get_command botnick msg pref = + let nl = String.length botnick in + let n = String.length pref in + let k = nl + 2 + n in + if String.get msg 0 = '!' then + (* If it's a command starting with "!" *) + let msg = String.sub msg 1 (String.length msg - 1) in + if String.length msg >= n then + Some (String.sub msg n (String.length msg - n)) + else None + else if String.length msg >= k && + String.sub msg 0 k = (botnick ^ ": " ^ pref) then + (* If someone is asking the bot directly, e.g. + vantabot: rollcall *) + Some (String.sub msg k (String.length msg - k)) + else None + + let callback_loop logf botnick connection result = + let open Irc_message in + match result with + | Result.Error e -> logf (formatf "Error parsing a message: %s" e) + | Result.Ok imsg -> + match imsg.command with + | PRIVMSG (target, msg) -> + if (List.mem target !bigchannels || List.mem target !smolchannels) + then + Lwt_list.iter_s + (fun (p,cb) -> + match get_command botnick msg p with + | None -> return () + | Some rest -> + cb connection (nick_of_irc_message imsg) target rest) + !callbacks + else return () + | _ -> return () + + + let pls_connect ~realname ~username ~server ~port ~nick ~logf ~onconnect = + C.set_log logf; + C.reconnect_loop + ~after:30 + ~connect:(fun () -> + logf "Connecting..." >>= fun () -> + C.connect_by_name ~realname ~username ~server ~port ~nick ()) + ~f:(fun connection -> + logf "Connected." >>= fun () -> onconnect connection) + ~callback:(callback_loop logf nick) + () + +end diff --git a/connection.mli b/connection.mli new file mode 100644 index 0000000..64cc1b1 --- /dev/null +++ b/connection.mli @@ -0,0 +1,38 @@ +open Irc_client +open Lwt + +(** We distinguish between "big channels" (allowed to spam) and "smol channels" (shouldn't spam). + Thanks to nilaky for the terminology. *) +module Conn (C: CLIENT with type 'a Io.t = 'a Lwt.t) : sig + type t = C.connection_t + + val smolchannels : string list ref + val bigchannels : string list ref + + val send_msg : t -> is_big:bool -> ?channel:string -> string -> unit Lwt.t + (** Send a message to all the channels associated with the connection. + set `is_big` (default: false) if the message is for `big channels` only. + If `channel` is specified then send to that channel only. *) + + val join_chan : t -> is_big:bool -> channel:string -> unit Lwt.t + (** Join a channel with and add it to the list. `is_big` default to false. *) + + val register_callback : + prefix:string -> + callback:(t -> string -> string -> string -> unit Lwt.t) -> + unit + (** Register a callback for the prefix "!". The arguments for + the callback are: connection nick (possible empty), channel, the rest of the + message. *) + + val pls_connect : + realname:string -> + username:string -> + server:string -> + port:int -> + nick:string -> + logf:(string -> unit Lwt.t) -> + onconnect:(t -> unit Lwt.t) -> + unit Lwt.t + +end diff --git a/vantabot.ml b/vantabot.ml index 4d6e28e..3731bc7 100644 --- a/vantabot.ml +++ b/vantabot.ml @@ -1,7 +1,8 @@ (* for tls support replace [Irc_client_lwt] with [Irc_client_tls] *) open Lwt open Irccolors -module C = Irc_client_lwt +open Connection +module C = Conn(Irc_client_lwt) let logf : string -> unit Lwt.t = Lwt_io.printl let formatf = Printf.sprintf @@ -13,24 +14,25 @@ let host = ref "localhost" let port = ref 6667 let realname = ":3" let botnick = ref "vantabot" -let channel = ref "#bots" +(* let channel = ref "#bots" *) let username = botnick let sleeptime = ref 30. let countdown = ref 5 let db_file = ref "toke.db" +let smolchannels = ["#tildetown"] +let bigchannels = ["#bots";"#aaa"] + let options = Arg.align [ "-host", Arg.Set_string host, " server host name (default: " ^ !host ^ ")" ; "-port", Arg.Set_int port, " server port (default: " ^ string_of_int !port ^ ")" - ; "-chan", Arg.Set_string channel, " target channel (default: " ^ !channel ^ ")" + (* ; "-chan", Arg.Set_string channel, " target channel (default: " ^ !channel ^ ")" *) ; "-nick", Arg.Set_string botnick, " nick prefix (default: " ^ !botnick ^ ")" ; "-sleeptime", Arg.Set_float sleeptime, " time period for joining the toke call, in seconds (default: " ^ string_of_float !sleeptime ^ ")" ; "-countdown", Arg.Set_int countdown, " countdown, in seconds (default: " ^ string_of_int !countdown ^ ")" ; "-db", Arg.Set_string db_file, " database file, will be created if does not exist (default: " ^ !db_file ^ ")" ] -let sendmsg connection m = C.send_privmsg ~connection ~target:!channel ~message:(color_text ~fg:Green ~bg:Black m) - (** * BUSINESS LOGIC HEH *) let toke_in_progress = ref false @@ -44,7 +46,7 @@ let push_toker t = let tokeplz () = let r = String.concat " " !tokers in - "SMOKE NOW " ^ r + bold_text (color_text ~fg:Green ~bg:Black "SMOKE NOW ") ^ r let inc_tokers_count () = List.iter (fun n -> Store.register_nick n (Store.DB !db_file)) !tokers @@ -60,29 +62,29 @@ let leaf = [ " WWW WWWWW WWW "; " WWW WWW WWW "; " WWW WWW WWW "; -" WWWWWWW "; +" WWWWWWW "; " WWWW | WWWW "; " | "; " | "] let show_leaf connection = - leaf |> Lwt_list.iter_s (fun l -> sendmsg connection l) + leaf |> Lwt_list.iter_s (fun l -> + C.send_msg connection ~is_big:true (bold_text (color_text ~fg:Green ~bg:Black l))) let start_toke connection = toke_in_progress := true; let rec toker_thread n conn = if n = 0 then begin - let msg = tokeplz () in + let msg = tokeplz () in inc_tokers_count (); tokers := []; - sendmsg connection msg - >>= fun () -> - toke_in_progress := false; + C.send_msg connection ~is_big:false msg >>= fun () -> + toke_in_progress := false; show_leaf connection end else begin - sendmsg connection (color_text ~bg:Black ~fg:White (formatf "ready in %i..." n)) + C.send_msg connection ~is_big:true (color_text ~bg:Green ~fg:LightGreen (formatf "get ready in %i..." n)) >>= fun () -> Lwt_unix.sleep 1. >>= fun () -> toker_thread (n-1) conn end @@ -91,84 +93,46 @@ let start_toke connection = let slt = !sleeptime in (if slt > 10. then Lwt_unix.sleep (slt-.10.) >>= fun () -> - sendmsg connection (color_text ~bg:Black ~fg:LightGreen "toke in ~10s") >>= fun () -> + C.send_msg connection ~is_big:true + (color_text ~bg:Black ~fg:LightGreen "get ready to toke in ~10s") >>= fun () -> Lwt_unix.sleep 10. else Lwt_unix.sleep slt) >>= fun () -> toker_thread !countdown connection); - sendmsg connection (formatf "Toke in progress, type in `!toke' to join. Time to go: ~%.1f secs" !sleeptime) + C.send_msg connection ~is_big:false (formatf "Toke in progress, type in `!toke' to join. Time to go: ~%.1f secs" !sleeptime) let handle_toke who conn = let is_new = push_toker who in if (!toke_in_progress) then if is_new - then sendmsg conn (formatf "%s joins the toke call" who) + then C.send_msg conn ~is_big:false (formatf "%s joins the toke call" who) else return () else start_toke conn -let rollcall ?who conn = +let rollcall who chan conn = let pref = match who with None -> "" | Some n -> n ^ ": " in - sendmsg conn (pref ^ "vantabot is a bot for synchronized toking. See . TL;DR: the only command is !toke") + C.send_msg conn ~channel:chan ~is_big:false (pref ^ "vantabot is a bot for synchronized toking. See . TL;DR: the only command is !toke") + +(* Register all the commands *) +let init conn = + C.register_callback "rollcall" (fun conn who chan _ -> + rollcall (if who = "" then None else Some who) chan conn); + C.register_callback "toke" (fun conn who _ _ -> + handle_toke (if who = "" then "" else who) conn) (***********************************) -(** * IRC connection handling code *) - - -let new_channel ~connection ~channel = - C.send_join ~connection ~channel - >>= fun () -> C.send_privmsg ~connection ~target:channel - ~message:(bold_text (color_text ~fg:Green ~bg:Black ".420.")) - -let nick_of_irc_message imsg = - let open Irc_message in - match imsg.prefix with - | None -> None - | Some wh -> - match String.split_on_char '!' wh with - | a::_ -> Some a - | [] -> None - - -let callback connection result = - let open Irc_message in - match result with - | Result.Error e -> logf (formatf "[callback] Error parsing a message: %s" e) - | Result.Ok imsg -> - logf (to_string imsg) >>= fun () -> - (match imsg.command with - | PRIVMSG (target, msg) -> - if target = !channel && String.length msg > 4 then - (* INITIATE OR JOIN THE TOKE *) - if String.sub msg 0 5 = "!toke" then - match nick_of_irc_message imsg with - | None -> return () - | Some wh -> handle_toke wh connection - (* RESPOND TO ROLLCALL *) - else if msg = (!botnick) ^ ": rollcall" then - match nick_of_irc_message imsg with - | None -> return () - | Some wh -> rollcall ~who:wh connection - else if msg = "!rollcall" then - rollcall connection - else return () - else return () - | _ -> return ()) - let lwt_main () = - C.set_log logf; - C.reconnect_loop - ~after:30 - ~connect:(fun () -> - logf "Connecting..." - >>= fun () -> C.connect_by_name ~realname:realname ~username:!username ~server:!host ~port:!port ~nick:!botnick () - ) - ~f:(fun connection -> - logf "Connected" - >>= fun () -> new_channel connection !channel) - ~callback - () + C.pls_connect ~realname:realname ~username:!username ~server:!host + ~port:!port ~nick:!botnick ~logf + ~onconnect:(fun connection -> + init connection; + Lwt_list.iter_p + (fun c -> C.join_chan connection ~is_big:false ~channel:c) smolchannels + >>= fun () -> + Lwt_list.iter_p + (fun c -> C.join_chan connection ~is_big:true ~channel:c) bigchannels) let _ = Arg.parse options (fun _ -> ()) " [options]\n";