From b1989bd79b31a10ce18cd98054aec95f52174873 Mon Sep 17 00:00:00 2001
From: Mylloon <kennel.anri@tutanota.com>
Date: Sat, 13 May 2023 17:10:19 +0200
Subject: [PATCH] add screenshot, wait for load, ability to run JS, add extra
 time in bonus so the twitter login page is loaded

---
 bin/main.ml |  5 +++--
 lib/json.ml |  7 +++++++
 lib/net.ml  | 37 +++++++++++++++++++++++++++++--------
 3 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/bin/main.ml b/bin/main.ml
index 773d0a4..528e72f 100644
--- a/bin/main.ml
+++ b/bin/main.ml
@@ -14,8 +14,9 @@ let stop (driver_process, session_id) =
 ;;
 
 let main session_id =
-  let body = navigate "https://twitter.com/i/flow/login" session_id in
-  print_endline body
+  ignore (navigate "https://twitter.com/i/flow/login" session_id);
+  (* Extra wait to be sure the page is loaded *)
+  Unix.sleep 5
 ;;
 
 let () =
diff --git a/lib/json.ml b/lib/json.ml
index fc4888a..6dc7304 100644
--- a/lib/json.ml
+++ b/lib/json.ml
@@ -19,3 +19,10 @@ let navigate_payload url = fmt {|
     "url": "%s"
   }
   |} url
+
+let execute_payload src = fmt {|
+  {
+    "script": "%s",
+    "args": []
+  }
+  |} src
diff --git a/lib/net.ml b/lib/net.ml
index fa0fe0a..0409003 100644
--- a/lib/net.ml
+++ b/lib/net.ml
@@ -2,6 +2,7 @@ open Cohttp_lwt_unix
 open Utils
 
 let driver_url = "http://127.0.0.1:4444/session"
+let driver id = fmt "%s/%s" driver_url id
 
 let send_post_request url json =
   let headers = Cohttp.Header.init_with "Content-Type" "application/json" in
@@ -42,14 +43,34 @@ let get_session () =
   | _ -> failwith "Invalid JSON"
 ;;
 
-let close_session id =
-  execute_delete_request (fmt "%s/%s" driver_url id) = "{\"value\":null}"
+let close_session id = execute_delete_request (driver id) = "{\"value\":null}"
+
+let execute_sync session_id src =
+  execute_post_request
+    (fmt "%s/execute/sync" (driver session_id))
+    (Json.execute_payload src)
 ;;
 
-let navigate url session_id =
-  ignore
-    (execute_post_request
-       (fmt "%s/%s/url" driver_url session_id)
-       (Json.navigate_payload url));
-  execute_get_request (fmt "%s/%s/title" driver_url session_id)
+let rec wait_for_load session_id =
+  let response = execute_sync session_id "return document.readyState" in
+  match Yojson.Safe.from_string response with
+  | `Assoc fields ->
+    (match List.assoc "value" fields with
+    | `String res ->
+      if not (res = "complete")
+      then (
+        Unix.sleep 1;
+        wait_for_load session_id)
+    | _ -> failwith "Error when waiting for load")
+  | _ -> failwith "Invalid JSON"
 ;;
+
+let navigate ?(wait = true) url session_id =
+  let res =
+    execute_post_request (fmt "%s/url" (driver session_id)) (Json.navigate_payload url)
+  in
+  if wait then wait_for_load session_id;
+  res
+;;
+
+let screenshot session_id = execute_get_request (fmt "%s/screenshot" (driver session_id))