#!/bin/bash

# This script is designed to create a livepatch module for the specified Photon OS kernel version.
#
# If GEN_LIVEPATCH_DEBUG is set to 1, the tmp directory will not be deleted at the end of the script. Useful for debugging purposes.
set -o pipefail
# shellcheck source=SPECS/kpatch/scripts/livepatch.sh
source ./livepatch.sh 2>/dev/null || source /usr/lib/livepatch.sh 2>/dev/null || {
    echo "Error: livepatch.sh not found in current directory or /usr/lib"
    exit 1
}
GCC=/usr/bin/gcc
BUILD_DIR="$(pwd)/gen_livepatch_workspace"
DEFAULT_OUTPUT_FOLDER="$(pwd)/output"
OUTPUT_FOLDER=""
TEMP_DIR="$BUILD_DIR/tmp"
RPM_BUILD_DIR="$TEMP_DIR/rpmbuild/BUILD"
VERSION_RELEASE_FLAVOR=""
LIVEPATCH_NAME=""
DEBUG_PKGNAME=""
KERNEL_SRC_RPM=""
SPEC_FILENAME=""
VMLINUX_PATH=""
PH_TAG=""
KERNEL_RELEASE=""
KERNEL_FLAVOR=""
KERNEL_VERSION=""
IS_RT=0
NON_REPLACE_FLAG=0
EXPORT_DEBUGINFO=0
DEBUGINFO_DIR=""
GENLOGS_DIR=""
DESC_FILE=""
RPM_DESCFILE=""
PACKAGE_AS_RPM=0
RPM_VERSION=
RPM_RELEASE=
RPM_NAME=""
BUILD_RPM_SPECFILE="/usr/share/livepatch/livepatch_spec.template"
patches=()
DEBUG_LOGS=0
SIGNC=""
SIGNC_UID=0
SIGNC_UNAME=""

# parse the command line arguments and fill in variables
module_name_string() {
    echo "${1//[^a-zA-Z0-9_-]/-}" | cut -c 1-55
}

parse_args() {
    # just print help message if no arguments
    if [ $# -eq 0 ]; then
        print_help 1
    elif [[ $1 != -* ]]; then
        echo "A flag must be set before any other parameters"
        print_help 1
    fi

    local opt=""
    local options=( "-s" "-p" "-v" "-o" "-h" "--help" "-k" "-n" "-R" "--export-debuginfo" "-d" "--rpm" "--rpm-version" "--rpm-release" "--rpm-desc" "--verbose" "--sign" "--signer-id" "--signer-name" )
    local no_arg_options=( "-R" "--export-debuginfo" "-h" "--help" "--rpm" "--verbose" "--sign" )
    while (( "$#" )); do
        if [[ $1 = -* ]]; then
            opt=$1

            # check to make sure number of args are correct
            if ! is_in_array "$opt" ${options[@]}; then
                error "Unknown option: $opt" >&2
            elif [[ $1 == -h || $1 == --help ]]; then
                print_help 0
            elif [[ $1 == -R ]]; then
                NON_REPLACE_FLAG=1
            elif [[ $1 == --export-debuginfo ]]; then
                EXPORT_DEBUGINFO=1
            elif [[ $1 == --rpm ]]; then
                PACKAGE_AS_RPM=1
            elif [[ $1 == --verbose ]]; then
                DEBUG_LOGS=1
            elif [[ $1 == --sign ]]; then
                SIGNC="/usr/bin/sign_livepatch"
                [[ -f "$SIGNC" ]] || error "Signc does not exist in the system: $SIGNC"
            elif ! is_in_array "$opt" ${no_arg_options[*]} && [[ $2 == -* || -z $2 ]]; then
                error "$1 needs at least one argument"
            elif  [[ $3 != -* && $opt != -p && -n $3 ]] && ! is_in_array "$opt" ${no_arg_options[*]} ; then
                error "$1 only takes one argument"
            fi
        else
            case "$opt" in
                -p)
                    patches+=("$1")
                    ;;
                -k)
                    VERSION_RELEASE_FLAVOR=$1
                    ;;
                -o)
                    OUTPUT_FOLDER=$1
                    ;;
                -n)
                    LIVEPATCH_NAME=$1
                    ;;
                -d)
                    DESC_FILE=$1
                    [[ -f "$DESC_FILE" ]] || error "Module description file does not exist"
                    DESC_FILE=$(readlink -f "$DESC_FILE")
                    ;;
                -s)
                    SRC_RPM_LOCAL_PATH=$1
                    [[ -f "$SRC_RPM_LOCAL_PATH" ]] || error "Unable to locate local src rpm at $SRC_RPM_LOCAL_PATH"
                    SRC_RPM_LOCAL_PATH=$(readlink -f "$SRC_RPM_LOCAL_PATH")
                    ;;
                -v)
                    DEBUGINFO_LOCAL_PATH=$1
                    [[ -f "$DEBUGINFO_LOCAL_PATH" ]] || error "Unable to locate local debuginfo rpm at $DEBUGINFO_LOCAL_PATH"
                    DEBUGINFO_LOCAL_PATH=$(readlink -f "$DEBUGINFO_LOCAL_PATH")
                    ;;
                --rpm-version)
                    RPM_VERSION=$1
                    ;;
                --rpm-release)
                    RPM_RELEASE=$1
                    ;;
                --rpm-desc)
                    RPM_DESCFILE=$1
                    [ -f "$RPM_DESCFILE" ] || error "RPM description file does not exist"
                    RPM_DESCFILE=$(readlink -f "$RPM_DESCFILE")
                    ;;
                --signer-id)
                    SIGNC_UID=$1
                    ;;
                --signer-name)
                    SIGNC_UNAME=$1
                    ;;
                esac
        fi
        # shift to the next argument
        shift
    done

    if [ -z "$patches" ]; then
        error "Please input at least one patch file"
    fi

    if [ -z "$VERSION_RELEASE_FLAVOR" ]; then
        echo "No kernel version specified, building native patch"
        VERSION_RELEASE_FLAVOR=$(uname -r)
    fi

    is_rt "$VERSION_RELEASE_FLAVOR"
    get_kernel_flavor "$VERSION_RELEASE_FLAVOR"
    get_kernel_version "$VERSION_RELEASE_FLAVOR"
    get_kernel_release "$VERSION_RELEASE_FLAVOR"
    get_photon_tag "$VERSION_RELEASE_FLAVOR"
    [[ -n "$SIGNC" ]] && config_sign_user

    if [[ "$PH_TAG" != *ph* ]]; then
        echo "Wrong kernel version detected: $KERNEL_VERSION_RELEASE"
        echo "Check for typos. Make sure it is in the same format as uname -r"
        exit 1
    fi

    [[ -z $RPM_VERSION ]] && RPM_VERSION=${KERNEL_VERSION}r${KERNEL_RELEASE}
    [[ -z $RPM_RELEASE ]] && RPM_RELEASE=0

    if [[ -z "$LIVEPATCH_NAME" ]]; then
        LIVEPATCH_NAME="linux-livepatch"
    elif [[ "${LIVEPATCH_NAME##*.}" == "ko" ]]; then
        LIVEPATCH_NAME="${LIVEPATCH_NAME%.*}"
    fi
    RPM_NAME="$LIVEPATCH_NAME"
    [[ $PACKAGE_AS_RPM -eq 1 ]] && LIVEPATCH_NAME="${LIVEPATCH_NAME}-${RPM_VERSION}-${RPM_RELEASE}"
    LIVEPATCH_NAME="$(module_name_string "$LIVEPATCH_NAME")"
    echo "Livepatch module name: $LIVEPATCH_NAME"

    if [ -z "$OUTPUT_FOLDER" ]; then
        echo "Output folder not specified, using default."
        OUTPUT_FOLDER="$DEFAULT_OUTPUT_FOLDER"
    fi
    DEBUGINFO_DIR="$OUTPUT_FOLDER"/debuginfo

    echo "Outputting livepatches to: $OUTPUT_FOLDER"
    GENLOGS_DIR="$OUTPUT_FOLDER/logs/"
    mkdir -p $GENLOGS_DIR
    GENLOGS_DIR=$(readlink -f "$GENLOGS_DIR")

    # if building an rpm, and rpm description is not specified, use the kernel module desc file if it is set
    if [[ -n "$DESC_FILE" && $PACKAGE_AS_RPM -eq 1 && -z "$RPM_DESCFILE" ]]; then
        echo "Separate RPM description not specified, using module description"
        RPM_DESCFILE=$DESC_FILE
    fi

    # make sure output folder exists
    mkdir -p "$OUTPUT_FOLDER"
}

config_sign_user() {
    echo "Adding user uid=$SIGNC_UID($SIGNC_UNAME) for signing."
    useradd -u $SIGNC_UID $SIGNC_UNAME 2>/dev/null || error "Cannot add signing user."
}

# takes in uname -r formatted string
# ex) 4.19.247-2.ph3-aws
get_kernel_flavor() {
    # check for rt kernel extension
    if [[ $IS_RT -ne 1 ]]; then
        KERNEL_FLAVOR=$(cut -d '-' -f 3 <<< "$1")
        if [ -z "$KERNEL_FLAVOR" ]; then
            KERNEL_FLAVOR="generic"
        fi
    else
        KERNEL_FLAVOR=rt
    fi
}

is_rt() {
    local rt_ext
    rt_ext="$(cut -d '-' -f 4 <<< "$1")"
    if [[ -n $rt_ext ]]; then
        IS_RT=1
    fi
}

# takes in uname -r formatted string
# returns the release number
get_kernel_release() {
    local field_num=2
    if [[ $IS_RT -eq 1 ]]; then
        # rt kernel
        field_num=3
    fi

    local release_tag release
    release_tag=$(cut -d '-' -f $field_num <<< "$1")
    release=$(cut -d '.' -f 1 <<< "$release_tag")
    KERNEL_RELEASE=$release
}

# takes in uname -r formatted string
# gets kernel version only
# ex) 5.10.118
get_kernel_version() {
    KERNEL_VERSION=$(cut -d '-' -f 1 <<< "$1")
}

# extracts .phX tag from
# uname -r formatted string
get_photon_tag() {
    local field_num=2
    if [[ $IS_RT -eq 1 ]]; then
        # rt kernel
        field_num=3
    fi

    local release_tag=$(cut -d '-' -f $field_num <<< "$1")
    local kernel_tag=$(cut -d '.' -f 2 <<< "$release_tag")
    PH_TAG=$kernel_tag
}

# vmlinux path is set when we parse the debuginfo package
set_kernel_package_paths() {
    #determine which photon version this is
    PHOTON_VERSION="${PH_TAG//[^0-9]/}".0

    #set file paths based on kernel flavor
    if [[ "$KERNEL_FLAVOR" == "generic" ]]; then
        DEBUG_PKGNAME="linux-debuginfo-$VERSION_RELEASE_FLAVOR.x86_64.rpm"
        KERNEL_SRC_RPM="linux-$VERSION_RELEASE_FLAVOR.src.rpm"
        SPEC_FILENAME="linux.spec"
    else
        DEBUG_PKGNAME="linux-$KERNEL_FLAVOR-debuginfo-$VERSION_RELEASE_FLAVOR.x86_64.rpm"
        KERNEL_SRC_RPM="linux-$KERNEL_FLAVOR-$VERSION_RELEASE_FLAVOR.src.rpm"
        SPEC_FILENAME="linux-$KERNEL_FLAVOR.spec"
    fi
}

print_config() {
    echo -e "\nBuilding livepatch"
    echo -e "Linux Version: $VERSION_RELEASE_FLAVOR"
    echo -e "Photon OS Version: $PHOTON_VERSION"
    echo -e "Photon OS Flavor: $KERNEL_FLAVOR"
    echo -e "Patch files: "
    for patch in "${patches[@]}"; do
        echo -e "\t $(basename "$patch")"
    done
}

# Copied from kpatch-build. Used to find gcc version that was used to compile an executable
gcc_version_from_file() {
    readelf -p .comment "$1" | grep -m 1 -o 'GCC:.*' || error "Error with readelf"
}

# Copied from kpatch-build. Used to check if GCC versions match
gcc_version_check() {
    local target="$1"
    local c="$BUILD_DIR/gcc_version_check.c"
    local o="$BUILD_DIR/gcc_version_check.o"
    local out gccver kgccver

    # gcc --version varies between distributions therefore extract version
    # by compiling a test file and compare it to vmlinux's version.
    echo 'int main(void) {return 0;}' > "$c"
    if ! out="$("$GCC" -c -pg -ffunction-sections -o "$o" "$c" 2>&1)" ; then
        error "GCC compilation error"
    fi

    if [[ -n "$out" ]]; then
        echo "gcc >= 4.8 required for -pg -ffunction-settings"
        echo "gcc output: $out"
        error
    fi

    gccver="$(gcc_version_from_file "$o")"
    kgccver="$(gcc_version_from_file "$target")"

    rm -f "$c" "$o"

    # ensure gcc version matches that used to build the kernel
    if [[ "$gccver" != "$kgccver" ]]; then
        echo "gcc/kernel version mismatch"
        echo "gcc version:    $gccver"
        echo "kernel version: $kgccver"
        echo "kpatch may have problems when building with the wrong gcc, exiting to be safe..."
        error
    fi
}

install_kernel_dependencies() {
    mapfile -t to_be_installed_pkgs < <(rpm -qpR "$1" | grep -vw rpmlib)
    echo "The following packages will be installed if they are not already installed:"
    echo "${to_be_installed_pkgs[*]}"
    tdnf install -qy "${to_be_installed_pkgs[@]}" || error "Error installing required packages"
}

parse_debuginfo_rpm() {
    echo -e "\nDownloading debug package and extracting vmlinux"

    if [[ -n "$DEBUGINFO_LOCAL_PATH" ]]; then
        cp "$DEBUGINFO_LOCAL_PATH" "$DEBUG_PKGNAME"
    else
        if ! wget --quiet --show-progress "https://packages.vmware.com/photon/$PHOTON_VERSION/photon_debuginfo_${PHOTON_VERSION}_x86_64/x86_64/$DEBUG_PKGNAME"; then
            error "Couldn't download photon kernel debug rpm"
        fi
    fi
    local absolute_path
    absolute_path=$(rpm -qlp "$DEBUG_PKGNAME" | grep "vmlinux-$KERNEL_VERSION")

    # remove the first slash from the path
    VMLINUX_PATH=${absolute_path:1}

    # extract vmlinux
    rpm2cpio "$DEBUG_PKGNAME" | cpio -ivd "./$VMLINUX_PATH"  > /dev/null || error "Couldn't extract vmlinux from $DEBUG_PKGNAME"
    cp "./$VMLINUX_PATH" "$TEMP_DIR"/rpmbuild/SOURCES/
}

download_source_rpm() {
    echo -e "\nDownloading source rpm, This will take a few minutes.\n"
    local src_rpm_url="https://packages.vmware.com/photon/${PHOTON_VERSION}/photon_srpms_${PHOTON_VERSION}_x86_64/$KERNEL_SRC_RPM"

    # allow downloading/copying of source rpm from either local or custom urls. Just need these variables to be exported before
    # running to enable these options.
    local KERNEL_SRC_RPM_PATH="$TEMP_DIR"/rpmbuild/SOURCES/"$KERNEL_SRC_RPM"
    if [ -n "${SRC_RPM_LOCAL_PATH}" ]; then
        cp "$SRC_RPM_LOCAL_PATH" "$KERNEL_SRC_RPM_PATH" || error "Couldn't find local src rpm"
    else
        if ! wget --quiet --show-progress "$src_rpm_url" -O "$KERNEL_SRC_RPM_PATH"; then
            error "Couldn't download photon kernel source rpm"
        fi
    fi

    install_kernel_dependencies "$KERNEL_SRC_RPM_PATH"


}

create_macro_file() {
    local linux_version
    [[ "$KERNEL_FLAVOR" == "generic" ]] && linux_version="$VERSION_RELEASE_FLAVOR" || linux_version="${VERSION_RELEASE_FLAVOR}-${KERNEL_FLAVOR}"

    local macrofile="$TEMP_DIR/rpmbuild/macros"
    echo "%define RPM_NAME $RPM_NAME" > $macrofile
    echo "%define dist .$PH_TAG" >> $macrofile
    echo "%define RPM_VERSION $RPM_VERSION" >> $macrofile
    echo "%define RPM_RELEASE $RPM_RELEASE" >> $macrofile
    echo "%define LINUX_VERSION $linux_version" >> $macrofile
    echo "%define KERNEL_SRC_RPM $KERNEL_SRC_RPM" >> $macrofile
    echo "%define LIVEPATCH_NAME $LIVEPATCH_NAME" >> $macrofile
    if [[ -n "$SIGNC" ]]; then
        echo "%define signing_script $SIGNC" >> $macrofile
        echo "%define signing_user $SIGNC_UNAME" >> $macrofile
    fi

    # For 3.0 kernels, default to using the -R flag since klp_replace is not supported.
    # For all other kernels, use the value specified by the user (if any).
    if [[ $NON_REPLACE_FLAG -eq 1 || $PHOTON_VERSION == "3.0" ]]; then
        echo "%define NON_REPLACE_FLAG -R" >> $macrofile
    else
        echo "%define NON_REPLACE_FLAG %{nil}" >> $macrofile
    fi

    if [[ $EXPORT_DEBUGINFO -eq 1 ]]; then
        echo "%define SKIP_CLEANUP --skip-cleanup" >> $macrofile
    else
        echo "%define SKIP_CLEANUP %{nil}" >> $macrofile
    fi

    #Copy patches and description files to sources directory
    cp "${patches[@]}" "$TEMP_DIR"/rpmbuild/SOURCES || error
    local sources_file="$TEMP_DIR"/rpmbuild/SOURCES/sources.inc
    >"$sources_file"

    local s_cnt=101
    echo -n "%define RPM_DESCRIPTION " >> $macrofile
    if [ -f "$RPM_DESCFILE" ]; then
        cp "$RPM_DESCFILE" "$TEMP_DIR"/rpmbuild/SOURCES
        echo "Source${s_cnt}: $(basename "$RPM_DESCFILE")" >>  "$sources_file"
        echo "%{expand:%(cat %{SOURCE${s_cnt}})}" >> $macrofile
        ((s_cnt++))
    else
        echo "Livepatch module for Linux $linux_version" >> $macrofile
    fi

    echo -n "%define PATCHES " >> $macrofile
    for patch in "${patches[@]}"; do
        echo "Source${s_cnt}: $(basename "$patch")" >>  "$sources_file"
        echo -n "%{SOURCE${s_cnt}} " >> $macrofile
        ((s_cnt++))
    done
    echo >> $macrofile

    if [[ -f "$DESC_FILE" ]]; then
        cp "$DESC_FILE" "$TEMP_DIR"/rpmbuild/SOURCES
        echo Source${s_cnt}: $(basename "$DESC_FILE") >>  "$sources_file"
        echo "%define DESCRIPTION_FILE -D %{SOURCE${s_cnt}}" >> $macrofile
        ((s_cnt++))
    else
        echo "%define DESCRIPTION_FILE %{nil}" >> $macrofile
    fi
}

# function to package the required module inside of an rpm
build_and_package_kpatch_module() {
    echo -e "\nPreparing to build livepatch"
    local spec_file="$TEMP_DIR/rpmbuild/SPECS/livepatch.spec"
    cp "$BUILD_RPM_SPECFILE" "$spec_file" || error
    echo "* $(date "+%a %b %d %Y") Auto Livepatch $RPM_VERSION-$RPM_RELEASE" >> "$spec_file"
    echo "- Auto Livepatch version $RPM_VERSION-$RPM_RELEASE built">> "$spec_file"

    create_macro_file

    if [[ $PACKAGE_AS_RPM -eq 1 ]]; then
       rpmbuild -ba $spec_file --define "%_topdir $TEMP_DIR/rpmbuild" &> $GENLOGS_DIR/rpmbuild-livepatch-pkg.log
    else
       rpmbuild -bb $spec_file --define "%_topdir $TEMP_DIR/rpmbuild" &> $GENLOGS_DIR/rpmbuild-livepatch-pkg.log
    fi
    status=$?
    [[ $DEBUG_LOGS -eq 1 || $status -ne 0 ]] && cat "$GENLOGS_DIR"/rpmbuild-livepatch-pkg.log
    [[ $status -ne 0 ]] && error "Building and/or packaging kernel module as rpm failed"
    echo "Livepatch module successfully built"

    # save all of the debug info like vmlinux, changed objects, etc if asked to
    # vmlinux is the patched vmlinux, not original
    if [[ $EXPORT_DEBUGINFO -eq 1 ]]; then
        echo "Saving kpatch-build debug files to $DEBUGINFO_DIR"

        [[ -n "$DEBUGINFO_DIR" ]] && rm -rf "${DEBUGINFO_DIR:?}"/*  > /dev/null
        mkdir -p "$DEBUGINFO_DIR"/patched-debuginfo
        cp -r ${HOME}/.kpatch/tmp "$DEBUGINFO_DIR"/kpatch-tmp
        cp ${HOME}/.kpatch/src/vmlinux* \
           ${HOME}/.kpatch/src/modules* \
           ${HOME}/.kpatch/src/Module.symvers \
           "$DEBUGINFO_DIR"/patched-debuginfo
    fi

    # should only build for x86_64 arch but put a * there just in case
    if [[ $PACKAGE_AS_RPM -eq 1 ]]; then
        cp "$TEMP_DIR"/rpmbuild/RPMS/*/"$RPM_NAME"*.rpm "$OUTPUT_FOLDER" || error "Current dir: $(pwd). Failed to save RPMs to $OUTPUT_FOLDER"
        cp "$TEMP_DIR"/rpmbuild/SRPMS/"$RPM_NAME"*.rpm "$OUTPUT_FOLDER" || error "Current dir: $(pwd). Failed to save source RPM to $OUTPUT_FOLDER"
    fi
    cp "$TEMP_DIR"/rpmbuild/BUILD/"$LIVEPATCH_NAME".ko "$OUTPUT_FOLDER"  || error "Current dir: $(pwd). Failed to save ko file to $OUTPUT_FOLDER"

    if [ -n "$SIGNC" ]; then
        while IFS= read -r -d '' file; do
            "$SIGNC" -t rpm -u "$SIGNC_UNAME" -a "$file" || error "Signing livepatch RPM failed!"
        done < <(find "$OUTPUT_FOLDER" -type f -name "*.rpm" -print0)
    fi
    echo "SUCCESS: Building rpm finished"
}

cleanup() {
    if [[ $GEN_LIVEPATCH_DEBUG -ne 1 ]]; then
        rm -rf $BUILD_DIR || error "Failed to delete temp dir: $BUILD_DIR"
    fi
}

#cleanup on exit
trap cleanup EXIT SIGINT SIGTERM

#make sure our working directories are clean before we start
rm -rf "$TEMP_DIR"
rm -rf "$GENLOGS_DIR"
rm -rf ${HOME}/.kpatch/

#parse command line arguments
parse_args "$@"

mkdir -p "$TEMP_DIR"
mkdir -p "$TEMP_DIR"/rpmbuild/{BUILD,BUILDROOT,RPMS,SOURCES,SPECS,SRPMS} || error

#set variables for all filenames/paths
set_kernel_package_paths

print_config

pushd $TEMP_DIR > /dev/null || error
#download and extract vmlinux
parse_debuginfo_rpm

gcc_version_check "$VMLINUX_PATH"

download_source_rpm
popd > /dev/null

echo -e "\nAll ready, building livepatch module."
echo -e "This may take up to one hour\n"
build_and_package_kpatch_module
