[edk2-staging/EdkRepo] [PATCH V1 2/3] EdkRepo: Add command completion setup to install.py

["Nate DeSimone" <nathaniel.l.desimone@intel.com>]

Add configuration of command completion scripts
to install.py This enables edkrepo command completions
to work "out of box" on most Linux systems by appending
to the user's .bashrc and .zshrc startup scripts
inclusion of the EdkRepo command completion scripts.

Cc: Ashley DeSimone <ashley.e.desimone@intel.com>
Cc: Puja Pandya <puja.pandya@intel.com>
Cc: Erik Bjorge <erik.c.bjorge@intel.com>
Cc: Prince Agyeman <prince.agyeman@intel.com>
Cc: Bret Barkelew <Bret.Barkelew@microsoft.com>
Cc: Philippe Mathieu-Daude <philmd@redhat.com>
Signed-off-by: Nate DeSimone <nathaniel.l.desimone@intel.com>
---
 edkrepo_installer/linux-scripts/install.py | 285 ++++++++++++++++++++-
 1 file changed, 283 insertions(+), 2 deletions(-)

diff --git a/edkrepo_installer/linux-scripts/install.py b/edkrepo_installer/linux-scripts/install.py
index b2cdfed..52f0c52 100755
--- a/edkrepo_installer/linux-scripts/install.py
+++ b/edkrepo_installer/linux-scripts/install.py
@@ -15,6 +15,7 @@ import importlib.util
 import logging
 import os
 import platform
+import re
 import stat
 import shutil
 import subprocess
@@ -22,7 +23,7 @@ import sys
 import traceback
 import xml.etree.ElementTree as et
 
-tool_sign_on = 'Installer for edkrepo version {}\nCopyright(c) Intel Corporation, 2019'
+tool_sign_on = 'Installer for edkrepo version {}\nCopyright(c) Intel Corporation, 2020'
 
 # Data here should be maintained in a configuration file
 cfg_dir = '.edkrepo'
@@ -31,6 +32,21 @@ cfg_src_dir = os.path.abspath('config')
 whl_src_dir = os.path.abspath('wheels')
 def_python = 'python3'
 
+# ZSH Configuration options
+prompt_regex = re.compile(r"#\s+[Aa][Dd][Dd]\s+[Ee][Dd][Kk][Rr][Ee][Pp][Oo]\s+&\s+[Gg][Ii][Tt]\s+[Tt][Oo]\s+[Tt][Hh][Ee]\s+[Pp][Rr][Oo][Mm][Pp][Tt]")
+zsh_autoload_compinit_regex = re.compile(r"autoload\s+-U\s+compinit")
+zsh_autoload_bashcompinit_regex = re.compile(r"autoload\s+-U\s+bashcompinit")
+zsh_autoload_colors_regex = re.compile(r"autoload\s+-U\s+colors")
+zsh_colors_regex = re.compile(r"\n\s*colors\n")
+zsh_compinit_regex = re.compile(r"compinit\s+-u")
+zsh_bashcompinit_regex = re.compile(r"\n\s*bashcompinit\n")
+zsh_autoload_compinit = 'autoload -U compinit'
+zsh_autoload_bashcompinit = 'autoload -U bashcompinit'
+zsh_autoload_colors = 'autoload -U colors'
+zsh_colors = 'colors'
+zsh_compinit = 'compinit -u'
+zsh_bashcompinit = 'bashcompinit'
+
 def default_run(cmd):
     return subprocess.run(cmd, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, check=True)
 
@@ -51,8 +67,85 @@ def get_args():
     parser.add_argument('-p', '--py', action='store', default=None, help='Specify the python command to use when installing')
     parser.add_argument('-u', '--user', action='store', default=None, help='Specify user account to install edkrepo support on')
     parser.add_argument('-v', '--verbose', action='store_true', default=False, help='Enables verbose output')
+    parser.add_argument('--no-prompt', action='store_true', default=False, help='Do NOT add EdkRepo combo and git branch to the shell prompt')
+    parser.add_argument('--prompt', action='store_true', default=False, help='Add EdkRepo combo and git branch to the shell prompt')
     return parser.parse_args()
 
+def is_prompt_customization_installed(user_home_dir):
+    script_files = [os.path.join(user_home_dir, '.bashrc'), os.path.join(user_home_dir, '.zshrc')]
+    customization_installed = True
+    for script_file in script_files:
+        if os.path.isfile(script_file):
+            with open(script_file, 'r') as f:
+                script = f.read().strip()
+        data = prompt_regex.search(script)
+        if not data:
+            customization_installed = False
+            break
+    return customization_installed
+
+__add_prompt_customization = None
+def get_add_prompt_customization(args, user_home_dir):
+    global __add_prompt_customization
+    if __add_prompt_customization is not None:
+        return __add_prompt_customization
+    if args.no_prompt:
+        __add_prompt_customization = False
+        return False
+    elif args.prompt:
+        __add_prompt_customization = True
+        return True
+    #Check if the prompt customization has already been installed
+    if is_prompt_customization_installed(user_home_dir):
+        __add_prompt_customization = False
+        return False
+    #If the prompt has not been installed and EdkRepo >= 2.0.0 is installed, then don't install the prompt customization
+    if shutil.which('edkrepo') is not None:
+        res = default_run(['edkrepo', '--version'])
+        if _check_version(res.stdout.replace('edkrepo ', '').strip(), '2.0.0') >= 0:
+            __add_prompt_customization = False
+            return False
+    #Show the user an advertisement to see if they want the prompt customization
+    from select import select
+    import termios
+    import tty
+    def get_key(timeout=-1):
+        key = None
+        old_settings = termios.tcgetattr(sys.stdin.fileno())
+        try:
+            tty.setraw(sys.stdin.fileno())
+            if timeout != -1:
+                rlist, _, _ = select([sys.stdin], [], [], timeout)
+                if rlist:
+                    key = sys.stdin.read(1)
+            else:
+                key = sys.stdin.read(1)
+        finally:
+            termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN, old_settings)
+        return key
+    print('\nEdkRepo can show the checked out \033[32mcombo\033[00m and \033[36mbranch\033[00m as part of the command prompt')
+    print('For example, instead of:\n')
+    print('\033[01;32muser@machine\033[00m:\033[01;34mEdk2\033[00m$ ')
+    print('\nThe command prompt would look like:\n')
+    print('\033[01;32muser@machine\033[00m:\033[01;34mEdk2\033[00m \033[32m[Edk2Master]\033[36m (master)\033[00m$ ')
+    print('')
+    while True:
+        print('Would you like the combo and branch shown on the command prompt? [y/N] ')
+        key = get_key(120)
+        if key:
+            if key == 'y' or key == 'Y':
+                print('Y')
+                __add_prompt_customization = True
+                return True
+            if key == 'n' or key == 'N':
+                print('N')
+                __add_prompt_customization = False
+                return False
+        else:
+            print('No response after 2min... assuming no.')
+            __add_prompt_customization = False
+            return False
+
 def get_installed_packages(python_command):
     pip_cmd = [def_python, '-m', 'pip', 'list', '--legacy']
     try:
@@ -146,6 +239,176 @@ def set_execute_permissions():
                     stat_data = os.stat(py_file)
                     os.chmod(py_file, stat_data.st_mode | stat.S_IEXEC)
 
+bash_prompt_customization = r'''
+# Add EdkRepo & git to the prompt
+ps1len="${#PS1}"
+let "pos3 = ps1len - 3"
+let "pos2 = ps1len - 2"
+if [ "${PS1:pos3}" == "\\$ " ]; then
+  newps1="${PS1:0:pos3}"
+  prompt_suffix="\\$ "
+elif [ "${PS1:pos3}" == " $ " ]; then
+  newps1="${PS1:0:pos3}"
+  prompt_suffix=" $ "
+elif [ "${PS1:pos2}" == "$ " ]; then
+  newps1="${PS1:0:pos2}"
+  prompt_suffix="$ "
+else
+  newps1="$PS1"
+  prompt_suffix=""
+fi
+
+# EdkRepo combo in prompt.
+if [ -x "$(command -v edkrepo)" ] && [ -x "$(command -v command_completion_edkrepo)" ]; then
+  newps1="$newps1\[\033[32m\]\$current_edkrepo_combo"
+  current_edkrepo_combo=$(command_completion_edkrepo current-combo)
+
+  # Determining the current Edkrepo combo requires invoking Python and parsing
+  # manifest XML, which is a relatively expensive operation to do every time
+  # the user presses <Enter>.
+  # As a performance optimization, only do this if the present working directory
+  # changed
+  if [[ ! -z ${PROMPT_COMMAND+x} ]] && [[ "$PROMPT_COMMAND" != "edkrepo_combo_chpwd" ]]; then
+    old_prompt_command=$PROMPT_COMMAND
+  fi
+  old_pwd=$(pwd)
+  edkrepo_combo_chpwd() {
+      if [[ "$(pwd)" != "$old_pwd" ]]; then
+        old_pwd=$(pwd)
+        current_edkrepo_combo=$(command_completion_edkrepo current-combo)
+      fi
+      if [[ ! -z ${PROMPT_COMMAND+x} ]]; then
+        eval $old_prompt_command
+      fi
+  }
+  PROMPT_COMMAND=edkrepo_combo_chpwd
+fi
+
+# Git branch in prompt.
+parse_git_branch() {
+  git branch 2> /dev/null | sed -e '/^[^*]/d' -e 's/* \(.*\)/ (\1)/'
+}
+
+PS1="$newps1\[\033[36m\]\$(parse_git_branch)\[\033[00m\]$prompt_suffix"
+'''
+
+zsh_prompt_customization = r'''
+# Add EdkRepo & git to the prompt
+prompt_length="${#PROMPT}"
+let "pos4 = prompt_length - 3"
+let "pos3 = prompt_length - 2"
+if [ "${PROMPT[$pos4,$prompt_length]}" = ' %# ' ]; then
+  new_prompt="${PROMPT[1,$pos4-1]}"
+  prompt_suffix=" %# "
+elif [ "${PROMPT[$pos3,$prompt_length]}" = "%# " ]; then
+  new_prompt="${PROMPT[1,$pos3-1]}"
+  prompt_suffix="%# "
+else
+  new_prompt="$PROMPT"
+  prompt_suffix=""
+fi
+
+# EdkRepo combo in prompt.
+if [ -x "$(command -v edkrepo)" ] && [ -x "$(command -v command_completion_edkrepo)" ]; then
+  new_prompt="$new_prompt%{$fg[green]%}\$current_edkrepo_combo%{$reset_color%}"
+  current_edkrepo_combo=$(command_completion_edkrepo current-combo)
+
+  # Determining the current Edkrepo combo requires invoking Python and parsing
+  # manifest XML, which is a relatively expensive operation to do every time
+  # the user presses <Enter>.
+  # As a performance optimization, only do this if the present working directory
+  # changed
+  function edkrepo_combo_chpwd() {
+    current_edkrepo_combo=$(command_completion_edkrepo current-combo)
+  }
+  chpwd_functions=(${chpwd_functions[@]} "edkrepo_combo_chpwd")
+fi
+
+# Load version control information
+autoload -Uz vcs_info
+precmd() { vcs_info }
+
+# Format the vcs_info_msg_0_ variable
+zstyle ':vcs_info:git:*' formats " %{$fg[cyan]%}(%b)%{$reset_color%}"
+
+# Set up the prompt (with git branch name)
+setopt PROMPT_SUBST
+eval "PROMPT='$new_prompt\${vcs_info_msg_0_}\$prompt_suffix'"
+'''
+
+def add_command_to_startup_script(script_file, regex, command, username):
+    script = ''
+    if os.path.isfile(script_file):
+        with open(script_file, 'r') as f:
+            script = f.read().strip()
+    data = regex.search(script)
+    if not data:
+        if script == '':
+            script = command
+        else:
+            script = '{}\n{}'.format(script, command)
+        with open(script_file, 'w') as f:
+            f.write(script)
+
+def add_command_comment_to_startup_script(script_file, regex, command, comment, username):
+    script = ''
+    if os.path.isfile(script_file):
+        with open(script_file, 'r') as f:
+            script = f.read().strip()
+    (new_script, subs) = re.subn(regex, command, script)
+    if subs == 0:
+        command = '\n{1}\n{0}\n'.format(command, comment)
+        if script == '':
+            new_script = command
+        else:
+            new_script = '{}\n{}'.format(script, command)
+    if new_script != script:
+        with open(script_file, 'w') as f:
+            f.write(new_script)
+        shutil.chown(script_file, user=username)
+        os.chmod(script_file, 0o644)
+
+def add_command_completions_to_shell(command_completion_script, args, username, user_home_dir):
+    # Add "source ~/.bashrc" to ~/.bash_profile if it does not have it already
+    bash_profile_file = os.path.join(user_home_dir, '.bash_profile')
+    bash_profile = ''
+    if os.path.isfile(bash_profile_file):
+        with open(bash_profile_file, 'r') as f:
+            bash_profile = f.read().strip()
+    profile_source_regex = re.compile(r"source\s+~/\.bashrc")
+    profile_source_regex2 = re.compile(r".\s+~/\.bashrc")
+    data = profile_source_regex.search(bash_profile)
+    if not data:
+        data = profile_source_regex2.search(bash_profile)
+        if not data:
+            if bash_profile == '':
+                bash_profile = 'source ~/.bashrc\n'
+            else:
+                bash_profile = '{}\nsource ~/.bashrc\n'.format(bash_profile)
+            with open(bash_profile_file, 'w') as f:
+                f.write(bash_profile)
+            shutil.chown(bash_profile_file, user=username)
+            os.chmod(bash_profile_file, 0o644)
+
+    # Add edkrepo command completion to ~/.bashrc if it does not have it already
+    regex = r"\[\[\s+-r\s+\"\S*edkrepo_completions.sh\"\s+\]\]\s+&&\s+.\s+\"\S*edkrepo_completions.sh\""
+    new_source_line = '[[ -r "{0}" ]] && . "{0}"'.format(command_completion_script)
+    comment = '\n# Add EdkRepo command completions'
+    bash_rc_file = os.path.join(user_home_dir, '.bashrc')
+    add_command_comment_to_startup_script(bash_rc_file, regex, new_source_line, comment, username)
+    if get_add_prompt_customization(args, user_home_dir):
+        add_command_to_startup_script(bash_rc_file, prompt_regex, bash_prompt_customization, username)
+    zsh_rc_file = os.path.join(user_home_dir, '.zshrc')
+    add_command_to_startup_script(zsh_rc_file, zsh_autoload_compinit_regex, zsh_autoload_compinit, username)
+    add_command_to_startup_script(zsh_rc_file, zsh_autoload_bashcompinit_regex, zsh_autoload_bashcompinit, username)
+    add_command_to_startup_script(zsh_rc_file, zsh_autoload_colors_regex, zsh_autoload_colors, username)
+    add_command_to_startup_script(zsh_rc_file, zsh_colors_regex, zsh_colors, username)
+    add_command_to_startup_script(zsh_rc_file, zsh_compinit_regex, zsh_compinit, username)
+    add_command_to_startup_script(zsh_rc_file, zsh_bashcompinit_regex, zsh_bashcompinit, username)
+    add_command_comment_to_startup_script(zsh_rc_file, regex, new_source_line, comment, username)
+    if get_add_prompt_customization(args, user_home_dir):
+        add_command_to_startup_script(zsh_rc_file, prompt_regex, zsh_prompt_customization, username)
+
 def do_install():
     global def_python
     org_python = None
@@ -209,6 +472,7 @@ def do_install():
         log.info('- Unable to determine users home directory')
         return 1
     default_cfg_dir = os.path.join(user_home_dir, cfg_dir)
+    get_add_prompt_customization(args, user_home_dir)
     log.info('+ System information collected')
 
     # Display current system information.
@@ -348,7 +612,7 @@ def do_install():
         #Delete obsolete dependencies
         if updating_edkrepo:
             installed_packages = get_installed_packages(def_python)
-            for whl_name in ['smmap2', 'gitdb2']:
+            for whl_name in ['smmap2', 'gitdb2', 'edkrepo-internal']:
                 if whl_name in installed_packages:
                     try:
                         res = default_run([def_python, '-m', 'pip', 'uninstall', '--yes', whl_name])
@@ -375,6 +639,23 @@ def do_install():
     set_execute_permissions()
     log.info('+ Marked scripts as executable')
 
+
+    #Install the command completion script
+    if shutil.which('edkrepo') is not None:
+        if args.local:
+            command_completion_script = os.path.join(default_cfg_dir, 'edkrepo_completions.sh')
+        else:
+            command_completion_script = os.path.join('/', 'etc', 'profile.d', 'edkrepo_completions.sh')
+        try:
+            res = default_run(['edkrepo', 'generate-command-completion-script', command_completion_script])
+            if args.local:
+                shutil.chown(command_completion_script, user=username)
+                os.chmod(command_completion_script, 0o644)
+            add_command_completions_to_shell(command_completion_script, args, username, user_home_dir)
+        except:
+            log.info('- Failed to configure edkrepo command completion')
+            if args.verbose:
+                traceback.print_exc()
     log.log(logging.PRINT, '\nInstallation complete\n')
 
     return 0
-- 
2.24.0.windows.2