#!/bin/bash

DOWNLOAD_ALL=""
CLEAR_CACHE=""
BUILD_BASE=""
DOWNLOAD_LANG=""
DOWNLOAD_VULN=""
CACHE_DIR="${HOME}/.cache/trivy-db/cache"
OUT_DIR="${HOME}/.cache/trivy-db/out"


function show_help(){
    echo -e "$(basename $0) [-h] [-d] [-c] [-a] [-b] [-l] [-v] [-t <cache dir>] [-o <out dir>]"
    echo -e "   -h Help"
    echo -e "   -d Debug mode"
    echo -e "   -с Clear cache directory"
    echo -e "   -a To download all advisories and other required files (${CACHE_DIR} dir by default)"
    echo -e "   -b Build database (${OUT_DIR} dir by default)"
    echo -e "   -l To download all other required files (${CACHE_DIR} dir by default)"
    echo -e "   -v To download all advisories (${CACHE_DIR} dir by default)"
    echo -e "   -t <cache dir> Dir for download all advisories"
    echo -e "   -o <out dir> Out dir for build database"
}

function get_lang(){
    CASHE_DIR=${1:-${CASHE_DIR}}
    mkdir -p ${CACHE_DIR}/{ruby-advisory-db,php-security-advisories,nodejs-security-wg,ghsa,cocoapods-specs,bitnami-vulndb,govulndb,k8s-cve-feed}
    echo "wget -qO - https://github.com/rubysec/ruby-advisory-db/archive/master.tar.gz | tar xz -C ${CACHE_DIR}/ruby-advisory-db --strip-components=1"
    wget -qO - https://github.com/rubysec/ruby-advisory-db/archive/master.tar.gz | tar xz -C ${CACHE_DIR}/ruby-advisory-db --strip-components=1
    echo "wget -qO - https://github.com/FriendsOfPHP/security-advisories/archive/master.tar.gz | tar xz -C ${CACHE_DIR}/php-security-advisories --strip-components=1"
    wget -qO - https://github.com/FriendsOfPHP/security-advisories/archive/master.tar.gz | tar xz -C ${CACHE_DIR}/php-security-advisories --strip-components=1
    echo "wget -qO - https://github.com/nodejs/security-wg/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/nodejs-security-wg --strip-components=1"
    wget -qO - https://github.com/nodejs/security-wg/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/nodejs-security-wg --strip-components=1
    echo "wget -qO - https://github.com/bitnami/vulndb/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/bitnami-vulndb --strip-components=1"
    wget -qO - https://github.com/bitnami/vulndb/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/bitnami-vulndb --strip-components=1
    echo "wget -qO - https://github.com/github/advisory-database/archive/refs/heads/main.tar.gz | tar xz -C ${CACHE_DIR}/ghsa --strip-components=1"
    wget -qO - https://github.com/github/advisory-database/archive/refs/heads/main.tar.gz | tar xz -C ${CACHE_DIR}/ghsa --strip-components=1
    echo "wget -qO - https://github.com/golang/vulndb/archive/refs/heads/master.tar.gz | tar xz -C ${CACHE_DIR}/govulndb --strip-components=1"
    wget -qO - https://github.com/golang/vulndb/archive/refs/heads/master.tar.gz | tar xz -C ${CACHE_DIR}/govulndb --strip-components=1
    ## required to convert GHSA Swift repo links to Cocoapods package names
    echo "wget -qO - https://github.com/CocoaPods/Specs/archive/master.tar.gz | tar xz -C ${CACHE_DIR}/cocoapods-specs --strip-components=1"
    wget -qO - https://github.com/CocoaPods/Specs/archive/master.tar.gz | tar xz -C ${CACHE_DIR}/cocoapods-specs --strip-components=1
    echo "wget -qO - https://github.com/kubernetes-sigs/cve-feed-osv/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/k8s-cve-feed --strip-components=1"
    wget -qO - https://github.com/kubernetes-sigs/cve-feed-osv/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/k8s-cve-feed --strip-components=1
}

function get_vuln(){
    CASHE_DIR=${1:-${CASHE_DIR}}
    REPO_OWNER=aquasecurity
    mkdir -p ${CACHE_DIR}/vuln-list
    echo "wget -qO - https://github.com/${REPO_OWNER}/vuln-list/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/vuln-list --strip-components=1"
    wget -qO - https://github.com/${REPO_OWNER}/vuln-list/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/vuln-list --strip-components=1
    mkdir -p ${CACHE_DIR}/vuln-list-redhat
    echo "wget -qO - https://github.com/${REPO_OWNER}/vuln-list-redhat/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/vuln-list-redhat --strip-components=1"
    wget -qO - https://github.com/${REPO_OWNER}/vuln-list-redhat/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/vuln-list-redhat --strip-components=1
    mkdir -p ${CACHE_DIR}/vuln-list-debian
    echo "wget -qO - https://github.com/${REPO_OWNER}/vuln-list-debian/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/vuln-list-debian --strip-components=1"
    wget -qO - https://github.com/${REPO_OWNER}/vuln-list-debian/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/vuln-list-debian --strip-components=1
    mkdir -p ${CACHE_DIR}/vuln-list-nvd
    echo "wget -qO - https://github.com/${REPO_OWNER}/vuln-list-nvd/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/vuln-list-nvd --strip-components=1"
    wget -qO - https://github.com/${REPO_OWNER}/vuln-list-nvd/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/vuln-list-nvd --strip-components=1
    mkdir -p ${CACHE_DIR}/vuln-list-aqua
    echo "wget -qO - https://github.com/${REPO_OWNER}/vuln-list-aqua/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/vuln-list-aqua --strip-components=1"
    wget -qO - https://github.com/${REPO_OWNER}/vuln-list-aqua/archive/main.tar.gz | tar xz -C ${CACHE_DIR}/vuln-list-aqua --strip-components=1
}

if [ $# -lt 1 ]
then
    show_help
    exit 1
fi

while getopts "dhcablvt:o:" opt; do
    case "$opt" in
    "d")
       set -x
    ;;
    "h")
        show_help
        exit 0
    ;;
    "a")
        DOWNLOAD_ALL="True"
    ;;
    "b")
        BUILD_BASE="True"
    ;;
    "c")
        CLEAR_CACHE="True"
    ;;
    "l")
        DOWNLOAD_LANG="True"
    ;;
    "v")
        DOWNLOAD_VULN="True"
    ;;
    "t")
        CACHE_DIR=$OPTARG
    ;;
    "o")
        OUT_DIR=$OPTARG
    ;;
    *)
        u=`basename $0`
        echo "$u: Unknown option: $1" >&2
        show_help
        exit 1
    ;;
    esac
done

if [[ -n ${CLEAR_CACHE} ]]
then
    echo "Clear cache"
    rm -rf ${CACHE_DIR}
    mkdir -pv ${CACHE_DIR}
fi

if [[ -n ${DOWNLOAD_ALL} ]]
then
    DOWNLOAD_LANG=True
    DOWNLOAD_VULN=True
fi

if [[ -n ${DOWNLOAD_LANG} ]]
then
    mkdir -pv ${CACHE_DIR}
    get_lang ${CACHE_DIR} || exit 1
fi

if [[ -n ${DOWNLOAD_VULN} ]]
then
    mkdir -pv ${CACHE_DIR}
    get_vuln ${CACHE_DIR} || exit 1
fi

if [[ -n ${BUILD_BASE} ]]
then
    if [[ ! -e ${CACHE_DIR} ]]
    then
        echo -e "Advisories and other required files directory ${CACHE_DIR} not found!\nTerminate..."
        echo -e "You must download advisories and other required files first!"
        show_help
        exit 1
    else
        mkdir -pv ${OUT_DIR}
        trivy-db build --cache-dir ${CACHE_DIR} --output-dir ${OUT_DIR} --update-interval 24h || exit 1
        bbolt compact -o ${OUT_DIR}/compact-trivy.db ${OUT_DIR}/trivy.db || exit 1
        rm -f ${OUT_DIR}/trivy.db || exit 1
        mv -f ${OUT_DIR}/compact-trivy.db ${OUT_DIR}/trivy.db || exit 1
        tar cvzf ${OUT_DIR}/db.tar.gz -C ${OUT_DIR} trivy.db metadata.json || exit 1
    fi
fi

