Merge branch 'dev' into 'master'

Dev

See merge request Mishio595/disml!3

6 files changed, 142 insertions, 60 deletions

diff --git a/bin/bot.ml b/bin/bot.ml
index c08218d..bf9f36e 100644
--- a/bin/bot.ml
+++ b/bin/bot.ml
@@ -1,30 +1,20 @@
 open Lwt.Infix
-open Animus
+open Disml
 
-let _ =
-    let data = Lwt_main.run (Http.get_gateway_bot ()) in
-    (* Yojson.Basic.pretty_print Format.std_formatter data;
-    print_newline (); *)
-    let url, _shards = match data with
-    | `Assoc [
-        ("url", `String url);
-        ("shards", `Int shards);
-        _
-    ] -> (url, shards)
-    | _ -> ("wss://gateway.discord.gg/", 1)
-    in
-    (Sharder.Shard.create {
-        url;
-        shards = [0; 1;];
-        token = Sys.getenv "DISCORD_TOKEN";
-    }
-    >|= fun (shard, recv_loop) ->
-    Lwt_engine.on_timer 60.0 true @@ begin
-    fun _ev -> Sharder.Shard.set_status shard
-    ("Current seq: " ^ string_of_int shard.seq)
-    >|= (fun _ -> print_endline "Status set!")
-    |> ignore;
+let main sharder =
+    Lwt_engine.on_timer 60.0 true begin
+        fun _ev -> Sharder.set_status_with sharder @@ begin
+        fun shard ->
+            `String ("Current seq: " ^ string_of_int shard.seq)
+        end
+        >|= (fun _ -> print_endline "Status set!")
+        |> ignore;
     end
+
+let _ =
+    Sharder.start @@ Sys.getenv "DISCORD_TOKEN"
+    >>= (fun sharder ->
+    main sharder
     |> ignore;
-    Lwt_main.run recv_loop)
+    sharder.promise)
     |> Lwt_main.run
\ No newline at end of file
diff --git a/dune b/dune
index 295e841..b8d2700 100644
--- a/dune
+++ b/dune
@@ -1,5 +1,5 @@
 (library
-	(name animus)
+	(name disml)
 	(modules endpoints http client sharder opcode)
 	(libraries lwt cohttp cohttp.lwt yojson websocket websocket-lwt-unix zlib)
 )
@@ -8,7 +8,7 @@
 (executable
     (name bot)
 	(modules bot)
-    (libraries lwt animus)
+    (libraries lwt disml)
 )
 
 (include_subdirs unqualified)
\ No newline at end of file
diff --git a/lib/client/client.ml b/lib/client/client.ml
index 3e90432..b1f6f25 100644
--- a/lib/client/client.ml
+++ b/lib/client/client.ml
@@ -1,5 +1,5 @@
-let notify _t _data =
-    (* Yojson.Basic.pretty_print Format.std_formatter data;
+let notify t data =
+    Yojson.Basic.pretty_print Format.std_formatter @@ `Assoc data;
     print_newline ();
-    print_endline t; *)
+    print_endline t;
     ()
\ No newline at end of file
diff --git a/lib/client/sharder.ml b/lib/client/sharder.ml
index 953a38f..18d1da0 100644
--- a/lib/client/sharder.ml
+++ b/lib/client/sharder.ml
@@ -1,6 +1,9 @@
 open Lwt.Infix
 open Websocket
 
+exception Invalid_Payload
+exception Invalid_Shards
+
 type data = {
     shards: int list;
     token: string;
@@ -11,7 +14,7 @@ module Shard = struct
     type t = {
         mutable hb: Lwt_engine.event option;
         mutable seq: int;
-        session: string option;
+        mutable session: string option;
         token: string;
         shard: int list;
         send: Frame.t -> unit Lwt.t;
@@ -19,6 +22,8 @@ module Shard = struct
         ready: unit Lwt.t;
     }
 
+    let id_rt = Lwt_mutex.create ()
+
     let parse (frame : Frame.t) =
         frame.content
         |> Yojson.Basic.from_string
@@ -40,7 +45,7 @@ module Shard = struct
         let frame = Frame.create ?content () in
         print_endline @@ Frame.show frame;
         shard.send frame
-        |> ignore;
+        >|= fun () ->
         shard
 
     let heartbeat shard =
@@ -59,26 +64,55 @@ module Shard = struct
         |> Yojson.Basic.Util.to_string in
         let seq = List.assoc "s" payload
         |> Yojson.Basic.Util.to_int in
-        let data = List.assoc "d" payload in
+        let data = List.assoc "d" payload
+        |> Yojson.Basic.Util.to_assoc in
         shard.seq <- seq;
         let _ = match t with
-        | "READY" -> Lwt.wakeup resolver ()
+        | "READY" ->
+            Lwt.wakeup resolver ();
+            let session = List.assoc "session_id" data
+            |> Yojson.Basic.Util.to_string in
+            shard.session <- Some session;
         | _ -> ()
         in
         Client.notify t data;
-        shard
+        Lwt.return shard
 
-    let set_status shard game =
+    let set_status shard status =
+        let payload = match status with
+        | `Assoc [("name", `String name); ("type", `Int t)] ->
+            `Assoc [
+                ("status", `String "online");
+                ("afk", `Bool false);
+                ("since", `Null);
+                ("game", `Assoc [
+                    ("name", `String name);
+                    ("type", `Int t)
+                ])
+            ]
+        | `String name ->
+            `Assoc [
+                ("status", `String "online");
+                ("afk", `Bool false);
+                ("since", `Null);
+                ("game", `Assoc [
+                    ("name", `String name);
+                    ("type", `Int 0)
+                ])
+            ]
+        | _ -> raise Invalid_Payload
+        in
+        shard.ready >|= fun _ -> push_frame ~payload shard STATUS_UPDATE
+
+    let request_guild_members ~guild ?(query="") ?(limit=0) shard =
         let payload = `Assoc [
-            ("status", `String "online");
-            ("afk", `Bool false);
-            ("since", `Null);
-            ("game", `Assoc [
-                ("name", `String game);
-                ("type", `Int 0)
-            ])
+            ("guild_id", `String (string_of_int guild));
+            ("query", `String query);
+            ("limit", `Int limit);
         ] in
-        shard.ready >|= fun _ -> push_frame ~payload shard STATUS_UPDATE
+        shard.ready >|= fun _ -> push_frame ~payload shard REQUEST_GUILD_MEMBERS
+
+
 
     let initialize shard data =
         print_endline "Initializing...";
@@ -96,14 +130,16 @@ module Shard = struct
         | Some s -> s
         in
         shard.hb <- Some hb;
+        Lwt_mutex.lock id_rt
+        >>= fun () ->
         match shard.session with
         | None ->
             let payload = `Assoc [
                 ("token", `String shard.token);
                 ("properties", `Assoc [
                     ("$os", `String Sys.os_type);
-                    ("$device", `String "animus");
-                    ("$browser", `String "animus")
+                    ("$device", `String "dis.ml");
+                    ("$browser", `String "dis.ml")
                 ]);
                 ("compress", `Bool false); (* TODO add compression handling*)
                 ("large_threshold", `Int 250);
@@ -116,7 +152,11 @@ module Shard = struct
                 ("session_id", `String s);
                 ("seq", `Int shard.seq)
             ] in
-            push_frame ~payload shard RECONNECT
+            push_frame ~payload shard RESUME
+        >|= fun s ->
+        Lwt_engine.on_timer 5.0 false (fun _ -> Lwt_mutex.unlock id_rt)
+        |> ignore;
+        s
 
     let handle_frame shard (term : Yojson.Basic.json) resolver =
         match term with
@@ -130,15 +170,17 @@ module Shard = struct
             match op with
             | DISPATCH -> dispatch shard term resolver
             | HEARTBEAT -> heartbeat shard
-            | RECONNECT -> print_endline "OP 7"; shard (* TODO reconnect *)
-            | INVALID_SESSION -> print_endline "OP 9"; shard (* TODO invalid session *)
-            | HELLO ->
-                let data = List.assoc "d" term in
-                initialize shard data
-            | HEARTBEAT_ACK -> shard
-            | opcode -> print_endline @@ "Invalid Opcode:" ^ Opcode.to_string opcode; shard
+            | RECONNECT -> print_endline "OP 7"; Lwt.return shard (* TODO reconnect *)
+            | INVALID_SESSION -> print_endline "OP 9"; Lwt.return shard (* TODO invalid session *)
+            | HELLO -> initialize shard @@ List.assoc "d" term
+            | HEARTBEAT_ACK -> Lwt.return shard
+            | opcode ->
+                print_endline @@ "Invalid Opcode:" ^ Opcode.to_string opcode;
+                Lwt.return shard
         end
-        | _ -> print_endline "Invalid payload"; shard
+        | _ ->
+            print_endline "Invalid payload";
+            Lwt.return shard
 
     let create data =
         let uri = (data.url ^ "?v=6&encoding=json") |> Uri.of_string in
@@ -158,7 +200,6 @@ module Shard = struct
             >>= fun frame ->
             let p = parse frame in
             handle_frame s p ready_resolver
-            |> Lwt.return
             >>= fun s -> recv_forever s
         end in
         let shard = {
@@ -174,6 +215,57 @@ module Shard = struct
         Lwt.return (shard, recv_forever shard)
 end
 
-type t = {
-    shards: Shard.t list;
-}
\ No newline at end of file
+type 'a t = {
+    shards: (Shard.t * 'a Lwt.t) list;
+    promise: 'a Lwt.t;
+}
+
+let start ?count token =
+    Http.get_gateway_bot ()
+    >|= fun data ->
+    let data = Yojson.Basic.Util.to_assoc data in
+    let url = List.assoc "url" data
+    |> Yojson.Basic.Util.to_string in
+    let count = match count with
+    | Some c -> c
+    | None -> List.assoc "shards" data
+        |> Yojson.Basic.Util.to_int
+    in
+    let shard_list = [0; count] in
+    let rec gen_shards l accum =
+        match l with
+        | [id; total;] when id < total ->
+            let shard_data = Lwt_main.run @@ Shard.create {
+                url;
+                shards = [id; total;];
+                token;
+            } in
+            shard_data :: gen_shards [id+1; total;] accum
+        | [_; _;] -> accum
+        | _ -> raise Invalid_Shards
+    in
+    let shards = gen_shards shard_list [] in
+    let p_list = List.map (fun (_, loop) -> loop) shards in
+    let promise = Lwt.choose p_list in
+    {
+        shards;
+        promise;
+    }
+
+let set_status sharder status =
+    List.map (fun (shard, _) ->
+        Shard.set_status shard status
+    ) sharder.shards
+    |> Lwt.nchoose
+
+let set_status_with sharder f =
+    List.map (fun (shard, _) ->
+        Shard.set_status shard @@ f shard
+    ) sharder.shards
+    |> Lwt.nchoose
+
+let request_guild_members ~guild ?query ?limit sharder =
+    List.map (fun (shard, _) ->
+        Shard.request_guild_members ~guild ?query ?limit shard
+    ) sharder.shards
+    |> Lwt.nchoose
\ No newline at end of file
diff --git a/lib/animus.ml b/lib/dis.ml
index e69de29..e69de29 100644
--- a/lib/animus.ml
+++ b/lib/dis.ml
diff --git a/lib/http.ml b/lib/http.ml
index 105ae34..b6b0298 100644
--- a/lib/http.ml
+++ b/lib/http.ml
@@ -20,7 +20,7 @@ module Base = struct
         let token = try
             Sys.getenv "DISCORD_TOKEN"
         with Not_found -> failwith "Please provide a token" in
-        let h = Header.init_with "User-Agent" "Animus v0.1.0" in
+        let h = Header.init_with "User-Agent" "Dis.ml v0.1.0" in
         let h = Header.add h "Authorization" ("Bot " ^ token) in
         Header.add h "Content-Type" "application/json"