Merge pull request #6 from endgameinc/support_cli_profiles

Support cli profiles

cli/main.ml

@@ -3,9 +3,9 @@ let sprintf = Printf.sprintf
 module R = Rresult.R
 open Engine.Lib
 
-let main repo_dir nocache deploy target_nodes dont_delete =
+let main repo_dir nocache deploy target_nodes dont_delete profile =
   let dont_delete = match dont_delete with | Some x -> [x] | None -> [] in
-  let params = Engine.Lib.make_params ~repo_dir ~nocache ~deploy ~target_nodes ~dont_delete in
+  let params = Engine.Lib.make_params ~repo_dir ~nocache ~deploy ~target_nodes ~dont_delete ?aws_profile:profile in
   Lwt_main.run (Engine.Runner.main params)
 
 let check repo_dir =
@@ -29,8 +29,8 @@ let check_configs_main cwd =
     (let%lwt configs = Engine.Runner.get_configs cwd in
      Lwt_list.iter_s Lwt_io.printl configs)
 
-let purge repo_dir target_hash () =
-    let%lwt creds = Aws_s3_lwt.Credentials.Helper.get_credentials () in
+let purge repo_dir target_hash profile () =
+    let%lwt creds = Aws_s3_lwt.Credentials.Helper.get_credentials ?profile () in
     let credentials = R.get_ok creds in
      let settings =
        Engine.Settings.parse_settings
@@ -49,9 +49,9 @@ let purge repo_dir target_hash () =
     let _ = Lwt_list.map_p check target_hash in
     Lwt.return ()
 
-let purge_main repo_dir target_hash =
+let purge_main repo_dir target_hash profile =
     Lwt_main.run (
-        purge repo_dir target_hash ()
+        purge repo_dir target_hash profile ()
     )
 
 let check_cache_per_node (node_name : string) cwd =
@@ -90,7 +90,7 @@ let dot_main repo_dir =
      in
      Lwt_io.printf "digraph { \n%s}\n" (String.concat "" content))
 
-let show_all_cache repo_dir =
+let show_all_cache profile repo_dir =
   Lwt_main.run
     (let%lwt () = Lwt_unix.chdir repo_dir in
      let%lwt new_configs = Engine.Runner.get_configs repo_dir in
@@ -117,7 +117,7 @@ let show_all_cache repo_dir =
      let print_node_cache (n : Engine.Node.real_node) =
        let name = Engine.Node.node_to_string (Engine.Node.Rnode n) in
        let%lwt rstatus =
-         Engine.Runner.AwsRunner.check_cache ~settings ~cwd:repo_dir ~n
+         Engine.Runner.AwsRunner.check_cache ?profile ~settings ~cwd:repo_dir ~n
        in
        let status = R.is_ok rstatus in
        let%lwt hash = Engine.Node.hash_of_node repo_dir n in
@@ -165,6 +165,10 @@ let dont_delete =
   let doc = "Don't delete a specific node, likely for debugging purposes." in
   Arg.(value & opt (some string) None & info ["dont-delete"] ~docv:"DONT_DELETE" ~doc)
 
+let aws_profile =
+  let doc = "Specific AWS profile to run under." in
+  Arg.(value & opt (some string) None & info ["profile"] ~docv:"AWS_PROFILE" ~doc)
+
 let make_info term_name doc =
   let man =
     [ `S Manpage.s_bugs
@@ -178,7 +182,7 @@ let invoke =
     make_info "invoke" "run the series of nodes to build a repository."
   in
   let term =
-    Term.(const main $ make_repo_dir 0 $ nocache $ deploy $ target_nodes $ dont_delete)
+    Term.(const main $ make_repo_dir 0 $ nocache $ deploy $ target_nodes $ dont_delete $ aws_profile)
   in
   (term, info)
 
@@ -219,15 +223,15 @@ let show_all_cache =
     make_info "show-all-cache"
       "print all nodes that can cache, their status, and current hash"
   in
-  let term = Term.(const show_all_cache $ make_repo_dir 0) in
+  let term = Term.(const show_all_cache $ aws_profile $ make_repo_dir 0) in
   (term, info)
 
 let purge_cache =
     let info =
         make_info "purge-cache"
         "purges one node's cached files for build"
     in
-    let term = Term.(const purge_main $ make_repo_dir 0 $ target_hash) in
+    let term = Term.(const purge_main $ make_repo_dir 0 $ target_hash $ aws_profile) in
     (term, info)
 
 let help =

dune-project

@@ -1,2 +1,2 @@
-(lang dune 1.3)
+(lang dune 2.7)
 (name makecloud)

engine/dune

@@ -1,6 +1,6 @@
 (library
  (name engine)
- (libraries jingoo lwt.unix cmdliner yaml.unix aws-ec2 aws-lwt aws-s3-lwt rresult uuidm digestif astring ppx_protocol_conv_yaml irmin-unix)
+ (libraries jingoo lwt.unix cmdliner yaml.unix aws-ec2 aws-lwt aws-s3-lwt rresult uuidm digestif astring ppx_protocol_conv_yaml ppx_protocol_conv irmin-unix )
  (preprocessor_deps "userdata_scripts/linux.txt" "userdata_scripts/windows.txt")
- (preprocess (pps lwt_ppx ppx_defer ocaml-monadic ppx_protocol_conv_yaml ppx_protocol_conv ppx_deriving.show ppx_deriving.eq ppx_blob ppx_irmin)))
+ (preprocess (pps lwt_ppx ppx_defer ocaml-monadic ppx_protocol_conv ppx_deriving.show ppx_deriving.eq ppx_blob ppx_irmin)))
 

engine/lib.ml

@@ -7,11 +7,12 @@ type run_parameters =
   ; deploy : bool
   ; target_nodes : string list
   ; dont_delete : string list
+  ; aws_profile : string option
   ; guid : Uuidm.t }
 
-let make_params ~repo_dir ~nocache ~deploy ~target_nodes ~dont_delete =
+let make_params ?aws_profile ~repo_dir ~nocache ~deploy ~target_nodes ~dont_delete =
   let guid = Uuidm.v4_gen (Random.State.make_self_init ()) () in
-  { repo_dir; nocache; deploy; target_nodes; guid; dont_delete }
+  { repo_dir; nocache; deploy; target_nodes; guid; dont_delete; aws_profile }
 
 let sprintf = Printf.sprintf
 

engine/provider_aws.ml

@@ -10,6 +10,7 @@ module Aws : Provider_template.Provider = struct
     ; settings: Settings.t
     ; aws_key: string
     ; aws_secret: string
+    ; aws_token: string option
     }
 
   let aws_to_result a =
@@ -38,7 +39,7 @@ module Aws : Provider_template.Provider = struct
     let b64 = Base64.(encode templated) |> R.failwith_error_msg in
     Lwt.return (String.split_on_char '=' b64 |> List.hd)
 
-  let apply_tags ~(settings : Settings.t) ~aws_key ~aws_secret ~guid ~instance_id () =
+  let apply_tags ~(settings : Settings.t) ~aws_key ~aws_secret ?token ~guid ~instance_id () =
     let tag_build = Types.Tag.make ~key:"Name" ~value:"Makecloud Builder" () in
     let tag_run = Types.Tag.make ~key:"Makecloud Run" ~value:guid () in
     let tags = [tag_build; tag_run] in
@@ -48,13 +49,14 @@ module Aws : Provider_template.Provider = struct
         ~region:settings.aws_region
         ~access_key:aws_key
         ~secret_key:aws_secret
+        ?token
         (module CreateTags)
         tags_req
     in
     Lwt.return (result |> aws_to_result)
 
-  let get_ami ~(settings : Settings.t) ~guid past_stage =
-    let%lwt creds = Aws_s3_lwt.Credentials.Helper.get_credentials () in
+  let get_ami ~(params : Lib.run_parameters) ~(settings : Settings.t) ~guid past_stage =
+    let%lwt creds = Aws_s3_lwt.Credentials.Helper.get_credentials ?profile:params.aws_profile () in
     let credentials = R.get_ok creds in
     let region = Aws_s3.Region.of_string settings.bucket_region in
     let endpoint =
@@ -69,23 +71,24 @@ module Aws : Provider_template.Provider = struct
     in
     Lwt.return @@ R.reword_error (fun _ -> `Msg "Failed to get ami in s3.") result
 
-  let get_base ~settings ~guid ~(n : Node.real_node) =
+  let get_base ~params ~settings ~guid ~(n : Node.real_node) =
     if String.sub n.base 0 (min (String.length n.base) 3) = "ami" then
       Lwt.return n.base
     else
-      let%lwt r = get_ami ~settings ~guid n.base in
+      let%lwt r = get_ami ~params ~settings ~guid n.base in
       Lwt.return (R.get_ok r)
 
-  let spinup (settings : Settings.t) (n : Node.real_node) guid : t Lwt.t =
+  let spinup (params : Lib.run_parameters) (settings : Settings.t) (n : Node.real_node) guid : t Lwt.t =
     let%lwt () = Node.node_log n "Node Spinning Up" in
     (* TODO improve this to new cred method? maybe pass creds around *)
-    let%lwt unsafe_creds = Aws_s3_lwt.Credentials.Helper.get_credentials () in
+    let%lwt unsafe_creds = Aws_s3_lwt.Credentials.Helper.get_credentials ?profile:params.aws_profile () in
     let creds = unsafe_creds |> ok_or_raise in
     let aws_key = creds.access_key in
     let aws_secret = creds.secret_key in
+    let aws_token = creds.token in
     let%lwt () = Node.node_log n "Got Credentials." in
     let%lwt user_data = make_user_data n settings in
-    let%lwt base = get_base ~settings ~guid ~n in
+    let%lwt base = get_base ~params ~settings ~guid ~n in
     let instance_params =
       (*TODO We should gen an ssh key, upload it to aws and use that instead of a constant key. *)
       (*TODO Pull instance size from the config file.*)
@@ -111,6 +114,7 @@ module Aws : Provider_template.Provider = struct
           ~region:settings.aws_region
           ~access_key:aws_key
           ~secret_key:aws_secret
+          ?token:aws_token
           (module RunInstances)
           instance_params
       in
@@ -127,7 +131,7 @@ module Aws : Provider_template.Provider = struct
     let get_details () =
       let result =
         Aws_lwt.Runtime.run_request ~region:settings.aws_region ~access_key:aws_key
-          ~secret_key:aws_secret
+          ~secret_key:aws_secret ?token:aws_token
           (module DescribeInstances)
           (Types.DescribeInstancesRequest.make ~instance_ids:[instance_id] ())
       in
@@ -154,7 +158,7 @@ module Aws : Provider_template.Provider = struct
     (*TODO aws_retry*)
     let%lwt ip = repeat_until_ok get_details 40 in
     let ip_address = R.failwith_error_msg ip in
-    Lwt.return {ip_address; instance_id; settings; aws_key; aws_secret}
+    Lwt.return {ip_address; instance_id; settings; aws_key; aws_secret; aws_token }
 
   let set_env t (n : Node.real_node) additional_env =
     let uri = Uri.of_string ("http://" ^ t.ip_address ^ ":8000/set_env") in
@@ -267,13 +271,14 @@ module Aws : Provider_template.Provider = struct
         ~region:settings.aws_region
         ~access_key:t.aws_key
         ~secret_key:t.aws_secret
+        ?token:t.aws_token
         (module CreateImage)
         image_req
     in
     Lwt.return (result |> aws_to_result)
 
-  let store_ami ~(settings : Settings.t) ~(n : Node.real_node) ~guid ami_id =
-    let%lwt creds = Aws_s3_lwt.Credentials.Helper.get_credentials () in
+  let store_ami ?profile ~(settings : Settings.t) ~(n : Node.real_node) ~guid ami_id =
+    let%lwt creds = Aws_s3_lwt.Credentials.Helper.get_credentials ?profile () in
     let credentials = R.get_ok creds in
     let region = Aws_s3.Region.of_string settings.bucket_region in
     let endpoint =
@@ -296,6 +301,7 @@ module Aws : Provider_template.Provider = struct
         ~region:settings.aws_region
         ~access_key:t.aws_key
         ~secret_key:t.aws_secret
+        ?token:t.aws_token
         (module DescribeImages)
         ami_req >>= fun x ->
         Lwt.return (aws_to_result x))
@@ -315,18 +321,18 @@ module Aws : Provider_template.Provider = struct
     | Error as e -> R.error_msgf "AMI has failed with reason %s."
       (Types.ImageState.to_string e)
 
-  let publish_image ~t ~settings ~n ~guid =
+  let publish_image ?profile ~t ~settings ~n ~guid =
     let instance_id = t.instance_id in
     let bind = Lwt_result.bind in
     let%bind box_id = repeat_until_ok (save_box ~t ~settings ~n ~instance_id ~guid) 20 in
     let none = (fun () -> R.error_msg missing_ami_err_msg) in
     let%bind image_id = Types.CreateImageResult.(box_id.image_id) |> R.of_option ~none |> Lwt.return
     in
     let%bind waiting_image = repeat_until_ok (check_on_ami ~t ~settings ~n image_id) 240 in
-    let%bind _store_result = store_ami ~settings ~n ~guid waiting_image in
+    let%bind _store_result = store_ami ?profile ~settings ~n ~guid waiting_image in
     Lwt.return_ok waiting_image
 
-  let runcmd transfer t (settings : Settings.t) (n : Node.real_node) guid (cmd : Command.t) :
+  let runcmd transfer t (params : Lib.run_parameters) (settings : Settings.t) (n : Node.real_node) guid (cmd : Command.t) :
       (string, [> R.msg] * string) result Lwt.t =
     let%lwt () = Node.node_log n (Command.to_string cmd) in
     let expire_time = 12 * Node.rnode_get_expire_time n in
@@ -340,7 +346,7 @@ module Aws : Provider_template.Provider = struct
       send_command t.ip_address ~expire_time
         @@ transfer ~first_arg ~second_arg ~verb:`Get
     | Publish ->
-      let%lwt image_id = publish_image ~t ~settings ~n ~guid in
+      let%lwt image_id = publish_image ?profile:params.aws_profile ~t ~settings ~n ~guid in
       (match image_id with
       | Ok i ->
         let%lwt () = Node.node_log n (Fmt.str "Saved instance as %s" i) in
@@ -352,7 +358,7 @@ module Aws : Provider_template.Provider = struct
     let%lwt () = Node.node_log n "Node spinning down." in
     let details () =
       Aws_lwt.Runtime.run_request ~region:t.settings.aws_region ~access_key:t.aws_key
-        ~secret_key:t.aws_secret
+        ~secret_key:t.aws_secret ?token:t.aws_token
         (module TerminateInstances)
         (Types.TerminateInstancesRequest.make ~instance_ids:[t.instance_id] ())
     in

engine/provider_stub.ml

@@ -3,7 +3,7 @@ open Lib
 module Stub : Provider_template.Provider = struct
   type t = unit
 
-  let spinup _settings (n : Node.real_node) _guid : t Lwt.t =
+  let spinup _params _settings (n : Node.real_node) _guid : t Lwt.t =
     let%lwt () =
       Lwt_io.printl ("[STUB]" ^ "[" ^ n.name ^ "]" ^ "Node Spinning Up")
     in
@@ -15,7 +15,7 @@ module Stub : Provider_template.Provider = struct
 
   let set_env _box _n _additional_env = Lwt.return ()
 
-  let runcmd _transfer_fn _settings _box (n : Node.real_node) _guid cmd =
+  let runcmd _transfer_fn _params _settings _box (n : Node.real_node) _guid cmd =
     let str_cmd = Command.to_string cmd in
     let _ = Lwt_io.printl ("[STUB]" ^ "[" ^ n.name ^ "]" ^ str_cmd) in
     Lwt.return (R.ok "")

engine/provider_template.ml

@@ -4,16 +4,16 @@ module type Provider = sig
   type t
 
   (*TODO: Investigate if we can bring Node.real_node into t*)
-  val spinup : Settings.t -> Node.real_node -> string -> t Lwt.t
+  val spinup : Lib.run_parameters -> Settings.t -> Node.real_node -> string -> t Lwt.t
 
   val set_env : t -> Node.real_node -> (string * string) list -> unit Lwt.t
 
   val wait_until_ready : t -> Node.real_node -> unit -> unit option Lwt.t
 
-  (*TODO: Commands need a type, really really need a type.*)
   val runcmd :
        (first_arg:string -> second_arg:string -> verb:verb -> string)
     -> t
+    -> Lib.run_parameters
     -> Settings.t
     -> Node.real_node
     -> string