Add MultipleProfile switching support

src/main/java/com/okta/tools/OktaAwsCliAssumeRole.java

@@ -26,6 +26,7 @@
 import com.amazonaws.services.securitytoken.model.AssumeRoleWithSAMLResult;
 import com.okta.tools.aws.settings.Configuration;
 import com.okta.tools.aws.settings.Credentials;
+import com.okta.tools.aws.settings.MultipleProfile;
 import com.okta.tools.saml.*;
 import org.apache.commons.lang.StringUtils;
 import org.apache.http.HttpStatus;
@@ -40,7 +41,6 @@
 import org.apache.http.message.BasicNameValuePair;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.ini4j.Profile;
 import org.json.JSONArray;
 import org.json.JSONException;
 import org.json.JSONObject;
@@ -75,7 +75,7 @@ static OktaAwsCliAssumeRole createOktaAwsCliAssumeRole(String oktaOrg, String ok
  return new OktaAwsCliAssumeRole(oktaOrg, oktaAWSAppURL, oktaAWSUsername, oktaAWSPassword,oktaProfile, oktaAWSRoleToAssume);
  }
 
- private OktaAwsCliAssumeRole(String oktaOrg, String oktaAWSAppURL, String oktaUsername, String oktaAWSPassword,String oktaProfile, String awsRoleToAssume) {
+ private OktaAwsCliAssumeRole(String oktaOrg, String oktaAWSAppURL, String oktaUsername, String oktaAWSPassword, String oktaProfile, String awsRoleToAssume) {
  this.oktaOrg = oktaOrg;
  this.oktaAWSAppURL = oktaAWSAppURL;
  this.oktaUsername = oktaUsername;
@@ -94,18 +94,28 @@ private static Path getSessionFilePath() {
 
  String run(Instant startInstant) throws Exception {
  Optional<Session> session = getCurrentSession();
- if (session.isPresent() && sessionIsActive(startInstant, session.get()))
- return session.get().profileName;
- String oktaSessionToken = getOktaSessionToken();
- String samlResponse = getSamlResponseForAws(oktaSessionToken);
- AssumeRoleWithSAMLRequest assumeRequest = chooseAwsRoleToAssume(samlResponse);
- Instant sessionExpiry = startInstant.plus(assumeRequest.getDurationSeconds() - 30, ChronoUnit.SECONDS);
- AssumeRoleWithSAMLResult assumeResult = assumeChosenAwsRole(assumeRequest);
- String profileName = createAwsProfile(assumeResult);
+ Path multiprofile = getMultipleProfilesIniPath();
+ FileReader reader = new FileReader(multiprofile.toFile());
+ MultipleProfile multipleProfile = new MultipleProfile(reader);
+ com.okta.tools.aws.settings.Profile profile = multipleProfile.getProfile(oktaProfile, multiprofile);
+
 
- updateConfigFile(profileName, assumeRequest.getRoleArn());
- updateCurrentSession(sessionExpiry, profileName);
- return profileName;
+ if (session.isPresent() && sessionIsActive(startInstant, session.get()) && oktaProfile.isEmpty())
+ return session.get().profileName;
+ if (profile == null || startInstant.isAfter(profile.expiry)) {
+ String oktaSessionToken = getOktaSessionToken();
+ String samlResponse = getSamlResponseForAws(oktaSessionToken);
+ AssumeRoleWithSAMLRequest assumeRequest = chooseAwsRoleToAssume(samlResponse);
+ Instant sessionExpiry = startInstant.plus(assumeRequest.getDurationSeconds() - 30, ChronoUnit.SECONDS);
+ AssumeRoleWithSAMLResult assumeResult = assumeChosenAwsRole(assumeRequest);
+ String profileName = createAwsProfile(assumeResult);
+
+ updateConfigFile(profileName, assumeRequest.getRoleArn());
+ addOrUpdateProfile(profileName, assumeRequest.getRoleArn(), sessionExpiry);
+ updateCurrentSession(sessionExpiry, profileName);
+ return profileName;
+ }
+ return oktaProfile;
  }
 
  private static class Session
@@ -119,12 +129,23 @@ private Session(String profileName, Instant expiry) {
  }
  }
 
- public static void logoutSession() throws IOException {
+ public void logoutSession() throws IOException {
+ if (oktaProfile != null) {
+ logoutMulti(oktaProfile);
+ }
  if (Files.exists(getSessionFilePath())) {
  Files.delete(getSessionFilePath());
  }
  }
 
+ private void logoutMulti(String oktaProfile) throws IOException {
+ Path multiprofile = getMultipleProfilesIniPath();
+ String profilestore = multiprofile.toString();
+ FileReader reader = new FileReader(multiprofile.toFile());
+ MultipleProfile multipleProfile = new MultipleProfile(reader);
+ multipleProfile.deleteProfile(profilestore,oktaProfile);
+ }
+
  private Optional<Session> getCurrentSession() throws IOException {
  if (Files.exists(getSessionFilePath())) {
  Properties properties = new Properties();
@@ -136,6 +157,7 @@ private Optional<Session> getCurrentSession() throws IOException {
  }
  return Optional.empty();
  }
+
  private boolean sessionIsActive(Instant startInstant, Session session) {
  return startInstant.isBefore(session.expiry);
  }
@@ -149,6 +171,19 @@ private String getOktaSessionToken() throws IOException {
  }
  }
 
+ private static Path getMultipleProfilesIniPath() throws IOException {
+ Path userHome = Paths.get(System.getProperty("user.home"));
+ Path oktaDir = userHome.resolve(".okta");
+ Path profileIni = oktaDir.resolve("profiles");
+ if (!Files.exists(oktaDir)) {
+ Files.createDirectory(oktaDir);
+ }
+ if(!Files.exists(profileIni)) {
+ Files.createFile(profileIni);
+ }
+ return profileIni;
+ }
+
  private void updateCurrentSession(Instant expiryInstant, String profileName) throws IOException {
  Properties properties = new Properties();
  properties.setProperty(OKTA_AWS_CLI_PROFILE_PROPERTY, profileName);
@@ -440,6 +475,17 @@ private void updateCredentialsFile(String profileName, String awsAccessKey, Stri
  }
  }
 
+ private void addOrUpdateProfile(String profileName,String oktaSession ,Instant start) throws IOException {
+ String profileStore = getMultipleProfilesIniPath().toString();
+ Reader reader = new FileReader(profileStore);
+ MultipleProfile multipleProfile = new MultipleProfile(reader);
+ multipleProfile.addOrUpdateProfile(profileName, oktaSession, start);
+ try (final FileWriter fileWriter = new FileWriter(profileStore)) {
+ multipleProfile.save(fileWriter);
+ }
+ }
+
+
  private void updateConfigFile(String profileName, String roleToAssume) throws IOException {
  String configLocation = System.getProperty("user.home") + "/.aws/config";
  boolean newConfiguration = !new File(configLocation).isFile();

src/main/java/com/okta/tools/WithOkta.java

@@ -18,7 +18,6 @@
  * limitations under the License.
  */
 
-
 public class WithOkta {
  public static void main(String[] args) throws Exception {
  OktaAwsConfig.createAwscli().run(Instant.now());
@@ -28,5 +27,4 @@ public static void main(String[] args) throws Exception {
  System.exit(exitCode);
  }
 
-}
-
+}

src/main/java/com/okta/tools/aws/settings/MultipleProfile.java

@@ -0,0 +1,91 @@
+/*
+ * Copyright 2017 Okta
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http:https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.okta.tools.aws.settings;
+
+import org.ini4j.Ini;
+import org.ini4j.Profile;
+
+import java.io.*;
+import java.nio.file.Path;
+import java.time.Instant;
+import java.util.Set;
+
+public class MultipleProfile extends Settings {
+
+ static final String SOURCE_PROFILE = "source_profile";
+ static final String PROFILE_EXPIRY = "profile_expiry";
+ static final String OKTA_SESSION = "okta_roleArn";
+
+ public com.okta.tools.aws.settings.Profile getProfile(String oktaprofile, Path profileIni) throws IOException {
+
+ Ini ini = new Ini(new File(profileIni.toString()));
+ Set<String> activeSessions = ini.keySet();
+ if (activeSessions.contains(oktaprofile)) {
+ Instant expiry = getExpiry(oktaprofile, ini);
+ String roleArn = getRoleArn(oktaprofile, ini);
+ return new com.okta.tools.aws.settings.Profile(expiry, roleArn);
+ }
+ return null;
+ }
+ public void deleteProfile(String profilestore, String oktaProfile) throws IOException {
+ try ( FileInputStream is = new FileInputStream(new File(profilestore))) {
+ Ini ini = new Ini(is);
+ ini.remove(ini.get(oktaProfile));
+ ini.store(new FileOutputStream(new File(profilestore)));
+ }
+ }
+ private Instant getExpiry(String oktaprofile,Ini ini) {
+ Ini.Section profilesection = ini.get(oktaprofile);
+ String profileExpiry = profilesection.get("profile_expiry");
+ Instant expiry = Instant.parse(profileExpiry);
+ return expiry;
+ }
+
+ private String getRoleArn(String oktaprofile, Ini ini) {
+ Ini.Section profilesection = ini.get(oktaprofile);
+ String roleArn = profilesection.get("okta_expiry");
+ return roleArn;
+ }
+
+ /**
+ * Create a Profiles object from a given {@link Reader}. The data given by this {@link Reader} should
+ * be INI-formatted.
+ *
+ * @param reader The settings we want to work with. N.B.: The reader is consumed by the constructor.
+ * @throws IOException Thrown when we cannot read or load from the given {@param reader}.
+ */
+ public MultipleProfile(Reader reader) throws IOException {
+ super(reader);
+ }
+
+ /**
+ * Add or update a profile to an Okta Profile file based on {@code name}. This will be linked to a okta profile
+ * of the same {@code name}, which should already be present in the profile expiry file.
+ * @param name The name of the profile.
+ * @param expiry expiry time of the profile session.
+ * @param okta_session the expiry time of the okta session
+ */
+ public void addOrUpdateProfile(String name, String okta_session,Instant expiry) {
+ final Profile.Section awsProfile = settings.get(name) != null ? settings.get(name) : settings.add(name);
+ writeSessionProfile(awsProfile,name,okta_session,expiry);
+ }
+
+ private void writeSessionProfile(Profile.Section awsProfile, String name, String okta_session, Instant expiry) {
+ awsProfile.put(SOURCE_PROFILE, name);
+ awsProfile.put(OKTA_SESSION, okta_session);
+ awsProfile.put(PROFILE_EXPIRY, expiry);
+ }
+}

src/main/java/com/okta/tools/aws/settings/Profile.java

@@ -0,0 +1,13 @@
+package com.okta.tools.aws.settings;
+
+import java.time.Instant;
+
+public class Profile {
+ public final Instant expiry;
+ public final String roleArn;
+
+ public Profile(Instant expiry, String roleArn) {
+ this.expiry = expiry;
+ this.roleArn = roleArn;
+ }
+}

src/main/java/com/okta/tools/awscli.java

@@ -21,13 +21,14 @@
 public class awscli {
  public static void main(String[] args) throws Exception {
  if (args.length > 0 && "logout".equals(args[0])) {
- OktaAwsCliAssumeRole.logoutSession();
+ OktaAwsCliAssumeRole oktaAwsCliAssumeRole = OktaAwsConfig.createAwscli();
+ oktaAwsCliAssumeRole.logoutSession();
  System.out.println("You have been logged out");
  System.exit(0);
  return;
  }
- String profileName = OktaAwsConfig.createAwscli().run(Instant.now());
  List<String> awsCommand = new ArrayList<>();
+ String profileName = OktaAwsConfig.createAwscli().run(Instant.now());
  awsCommand.add("aws");
  awsCommand.add("--profile");
  awsCommand.add(profileName);