#!/bin/sh
logFile="/var/log/install_virtio_drvs_to_initramfs.log"
log () {
    date=`date +"%Y-%m-%d %H:%M:%S"`
    echo "[$date] $1" | tee -a $logFile
}

kerVer=$(uname -r)
dt=`date +%m-%d-%H%M%S`

virtioBlkFile1="/lib/modules/$kerVer/kernel/drivers/block/virtio_blk.ko"
virtioBlkFile2="/lib/modules/$kerVer/kernel/drivers/block/virtio_blk.ko.xz"
if [ ! -f "$virtioBlkFile1" ] && [ ! -f "$virtioBlkFile2" ]; then
	log "not find virtio blk drivers in system."
	log "build virtio drivers depend on the follow web page:"
	log "https://www.alibabacloud.com/help/zh/doc-detail/62423.htm"
	log "after build virtio drivers, execute this script again."
	log "exit."
	exit -1
fi 

rhInitrdPath="/boot/initramfs-$kerVer.img"
udInitrdPath="/boot/initrd.img-$kerVer"
suseInitrdPath="/boot/initrd-$kerVer"

if [ -f $rhInitrdPath ]; then
	log "system based on centos/redhat, try to install virtio drivers to initramfs."

	output=`lsinitrd $rhInitrdPath 2>&1`
	if [ $? != 0 ]; then
		log $output
		log "check virtio drivers installation failed, exit."
		exit -1
	fi
	if [[ $output == *"virtio_blk"* ]]; then
		log "virtio driver already installed in initramfs, exit."
		exit 0
	fi

	log "backup initramfs to $rhInitrdPath.$dt.bak"
	output=`cp $rhInitrdPath $rhInitrdPath.$dt.bak 2>&1`
	if [ $? != 0 ]; then
		log "backup initramfs failed:"
		log $output
		log "exit."
		exit -1
	fi

	log "start to install virtio drivers to initramfs."
	output=`mkinitrd -f --allow-missing \
            --with=xen-blkfront --preload=xen-blkfront \
            --with=virtio_blk --preload=virtio_blk \
            --with=virtio_pci --preload=virtio_pci \
            --with=virtio_console --preload=virtio_console \
            $rhInitrdPath $kerVer 2>&1`
	ret=$?
	log $output
	if [ $ret != 0 ]; then
		log "install virtio drivers to initramfs failed, recover from backup."
		mv $rhInitrdPath $rhInitrdPath.failed
		cp $rhInitrdPath.$dt.bak $rhInitrdPath
		log "exit."
		exit -1
	else
		log "install virtio drivers to initramfs successfully."
		exit 0
	fi
fi

if [ -f $udInitrdPath ]; then
	log "system based on ubuntu/debian, try to install virtio drivers to initrd."

	output=`lsinitramfs $udInitrdPath 2>&1`
	if [ $? != 0 ]; then
		log $output
		log "check virtio drivers installation failed, exit."
		exit -1
	fi
	if [[ $output == *"virtio_blk"* ]]; then
		log "virtio drivers already installed in initrd, exit."
		exit 0
	fi

	log "backup initrd to $udInitrdPath.$dt.bak"
	output=`cp $udInitrdPath $udInitrdPath.$dt.bak 2>&1`
	if [ $? != 0 ]; then
		log "backup initrd failed:"
		log $output
		log "exit."
		exit -1
	fi

	log "start to install virtio drivers to initrd."
	echo -e 'xen-blkfront\nvirtio_blk\nvirtio_pci\nvirtio_console' >> /etc/initramfs-tools/modules
    output=`mkinitramfs -o $udInitrdPath 2>&1`

	ret=$?
	log $output
	if [ $ret != 0 ]; then
		log "install virtio drivers to initrd failed, recover from backup."
		mv $udInitrdPath $udInitrdPath.failed
		cp $udInitrdPath.$dt.bak $udInitrdPath
		log "exit."
		exit -1
	else
		log "install virtio drivers to initrd successfully."
		exit 0
	fi
fi

if [ -f $suseInitrdPath ]; then
	log "system based on suse, try to install virtio drivers to initrd."

	output=`lsinitrd $suseInitrdPath 2>&1`
	if [ $? != 0 ]; then
		log $output
		log "check virtio drivers installation failed, exit."
		exit -1
	fi
	if [[ $output == *"virtio_blk"* ]]; then
		log "virtio driver already installed in initrd, exit."
		exit 0
	fi

	log "backup initrd to $suseInitrdPath.$dt.bak"
	output=`cp $suseInitrdPath $suseInitrdPath.$dt.bak 2>&1`
	if [ $? != 0 ]; then
		log "backup initrd failed:"
		log $output
		log "exit."
		exit -1
	fi

	log "start to install virtio drivers to initrd."
	output=`dracut --add-drivers "ata_piix ata_generic virtio_blk \
            virtio_scsi virtio_net virtio_pci virtio_ring virtio" \
            --force 2>&1`
	ret=$?
	log $output
	if [ $ret != 0 ]; then
		log "install virtio drivers to initrd failed, recover from backup."
		mv $suseInitrdPath $suseInitrdPath.failed
		cp $suseInitrdPath.$dt.bak $suseInitrdPath
		log "exit."
		exit -1
	else
		log "install virtio drivers to initrd successfully."
		exit 0
	fi
fi
