Move cog setup logic into shell_utils so it can be used when any of these commands are run:

- `m`
- `lunch`
- `source build/envsetup.sh`

this will set up the symlink properly regardless of the order of commands so users don't get stuck with a broken build

Bug: 362337892
Change-Id: Ibd3b262107dfc1024cd83ab8aeb33a299fb0ffb2
diff --git a/shell_utils.sh b/shell_utils.sh
index 86f3f49..6cc8c9e 100644
--- a/shell_utils.sh
+++ b/shell_utils.sh
@@ -63,6 +63,75 @@
 }
 fi
 
+# This function sets up the build environment to be appropriate for Cog.
+function setup_cog_env_if_needed() {
+  local top=$(gettop)
+
+  # return early if not in a cog workspace
+  if [[ ! "$top" =~ ^/google/cog ]]; then
+    return 0
+  fi
+
+  setup_cog_symlink
+
+  export ANDROID_BUILD_ENVIRONMENT_CONFIG="googler-cog"
+
+  # Running repo command within Cog workspaces is not supported, so override
+  # it with this function. If the user is running repo within a Cog workspace,
+  # we'll fail with an error, otherwise, we run the original repo command with
+  # the given args.
+  if ! ORIG_REPO_PATH=`which repo`; then
+    return 0
+  fi
+  function repo {
+    if [[ "${PWD}" == /google/cog/* ]]; then
+      echo -e "\e[01;31mERROR:\e[0mrepo command is disallowed within Cog workspaces."
+      kill -INT $$ # exits the script without exiting the user's shell
+    fi
+    ${ORIG_REPO_PATH} "$@"
+  }
+}
+
+# creates a symlink for the out/ dir when inside a cog workspace.
+function setup_cog_symlink() {
+  local out_dir=$(getoutdir)
+  local top=$(gettop)
+
+  # return early if out dir is already a symlink
+  if [[ -L "$out_dir" ]]; then
+    return 0
+  fi
+
+  # return early if out dir is not in the workspace
+  if [[ ! "$out_dir" =~ ^$top/ ]]; then
+    return 0
+  fi
+
+  local link_destination="${HOME}/.cog/android-build-out"
+
+  # if there's a local out/ dir, prompt the user to remove it.
+  # fail out if they say no.
+  local answer
+  if [[ -d "$out_dir" ]]; then
+    echo "Detected existing out/ directory in the Cog workspace which is not supported. Can we repair your workspace by removing it and creating the symlink to ~/.cog/android-build-out instead? (y/N): "
+    read -r answer
+    if [[ $answer =~ ^[Yy]$ ]]; then
+      rm -rf "$out_dir"
+    else
+      echo "Exiting due to unsupported out/ directory."
+      kill -INT $$ # exits the script without exiting the user's shell
+    fi
+  fi
+
+  # create symlink
+  echo "Creating symlink: $out_dir -> $link_destination"
+  mkdir -p ${link_destination}
+  if ! ln -s "$link_destination" "$out_dir"; then
+    echo "Failed to create cog symlink: $out_dir -> $link_destination" >&2
+    kill -INT $$ # exits the script without exiting the user's shell
+  fi
+}
+
 function getoutdir
 {
     local top=$(gettop)