def cb_register_user_gcs()

in community/front-end/ofe/infrastructure_files/gcs_bucket/clusters/ansible_setup/roles/c2_daemon/files/ghpcfe_c2daemon.py [0:0]

• 93 lines of code
• 10 McCabe index (conditional complexity)

def cb_register_user_gcs(message, **kwargs):
    """Handle registration of user GCS credentials"""
    if not "ackid" in message:
        logger.error(
            "Refusing REGISTER_USER_GCS message without ackid (message was %s)",
            message,
        )
        return
    ackid = message["ackid"]
    response = {"ackid": ackid}

    logger.info("Starting REGISTER_USER_GCS: %s", message)

    try:
        (username, unused_uid, unused_gid, homedir) = _verify_oslogin_user(
            message["login_uid"]
        )
    except KeyError:
        logger.error(
            "User with uid=%s not OS-Login enabled", message["login_uid"]
        )
        response["status"] = "User does not have OS-Login permissions"
        response[
            "message"
        ] = "User is not allowed to submit jobs to this cluster"
        send_message("ACK", response)
        return

    try:
        response["status"] = "Configuring gcloud"
        send_message("UPDATE", response)
        subprocess.run(
            [
                "sudo",
                "-u",
                username,
                "gcloud",
                "config",
                "set",
                "pass_credentials_to_gsutil",
                "false",
            ],
            check=True,
        )

        # gsutil will fail if the backup file already exists
        boto_backup = Path(homedir) / ".boto.bak"
        if boto_backup.exists():
            boto_backup.unlink()

        with pexpect.spawn(
            "sudo",
            args=[
                "-u",
                username,
                "gsutil",
                "config",
                "-s",
                "https://www.googleapis.com/auth/devstorage.read_write",
            ],
        ) as child:
            child.expect(
                "Please navigate your browser to the following script_url:"
            )
            child.readline()  # Eat newline
            url = str(child.readline(), "utf-8").strip()
            response["status"] = "Waiting For User Auth"
            response["verify_url"] = url

            # Set up wait signal
            my_verify_key = None

            def my_callback(message):
                nonlocal my_verify_key
                my_verify_key = message.get("verify_key", None)

            _c2_ackMap[ackid] = my_callback

            send_message("UPDATE", response)
            response.pop("verify_url")

            # Wait for user to auth
            attempts = 0
            while not my_verify_key:
                time.sleep(2)
                attempts += 1
                if attempts > 150:  # 300 seconds
                    logger.error("Wait timed out - 5 minutes passed!")
                    response["status"] = "Wait timed out - 5 minutes passed!"
                    send_message("ACK", response)
                    child.terminate(force=True)
                    _c2_ackMap.pop(ackid)
                    return

            # Remove our callback, now that we have our verify key
            _c2_ackMap.pop(ackid)

            child.expect("Enter the authorization code:")
            child.sendline(my_verify_key)
            child.expect(pexpect.EOF)
            child.wait()
            child.close()
            response["exit_status"] = child.exitstatus
            response["status"] = (
                "Success" if child.exitstatus == 0 else "Failure"
            )
            send_message("ACK", response)

    except Exception as err:
        logger.error("Failed to configure User's GCS creds.", exc_info=err)
        send_message("ACK", response)