#!/bin/bash
# License: GPL 
# Author: Steven Shiau <steven _at_ clonezilla org>
# Description: Program to start BT service for restoring
# Including 3 parts:
# Part 1: create metainfo file (.torrent)
# Part 2: start tracker
# Part 3: start bitorrent client program

#
DRBL_SCRIPT_PATH="${DRBL_SCRIPT_PATH:-/usr/share/drbl}"
. $DRBL_SCRIPT_PATH/sbin/drbl-conf-functions
. /etc/drbl/drbl-ocs.conf
. $DRBL_SCRIPT_PATH/sbin/ocs-functions

# Load the config in ocs-live.conf. This is specially for Clonezilla live. It will overwrite some settings of /etc/drbl/drbl-ocs.conf, such as $DIA...
[ -e "/etc/ocs/ocs-live.conf" ] && . /etc/ocs/ocs-live.conf

# Settings
# metainfo_creator and bt_client is default defined in drbl-ocs.conf, or command option -t|--bt-client.
# Flag to brutally kill the related service
brutal_kill="no"
# BT_PORT_INIT, p_length_ctorrent, p_length_mktorrent, ocsroot_btzone, btlog_dir are loaded from drbl-ocs.conf.
# Log file to keep the status of ezio seeding.
ezio_seeding_log="/var/log/ezio-seeding"
# Default mode for reading BT data from image
bt_slices_from_dev="img"

# Functions
USAGE() {
    echo "$ocs - To start the bittorrent for image dir"
    echo "Usage:"
    echo "To run $ocs:"
    echo "$ocs [OPTION] [start|stop|status] [IMAGE] [DEV1] [DEV1]..."
    echo "Options:"
    echo "-b, --brutal-kill            Brutally kill all the related service when stopping."
    echo "-c, --metainfo-creator PROG  Use PROG to create the metainfo (.torrent) file"
    echo "-f, --from-raw-dev           Force to read BT data from raw device. Default is from image"
    echo "-g, --log-dir DIR            Use DIR as the log dir. Default value is \"$btlog_dir\""
    echo "-m, --max-client-no NO       Assign the client number as NO. When this number of clients finishes downloading, BT tracker will exit so that client will stop, too."
    echo "-p, --bt-srv-ip  IP          Assign the IP address for using in tracker's URL. If not assigned, the IP address detected on the machine will be used"
    echo "-r, --bt-root DIR            Use DIR as the working dir. It's preceded to the files downloaded. Default value is \"$ocsroot_btzone\""
    echo "-t, --bt-client PROG         Use PROG instead of ctorrent as the bittorrent client program. Available PROG: ezio, ctorrent, lftp, btdownloadheadless, aria2c"
    echo "IMAGE is the image name under \"$ocsroot_btzone\""
    echo "DEV* is the device name"
    echo "Ex:"
    echo "To start the bittorrent service for the image dir \"/home/partimag/btzone/myimg/sda1/\" and \"/home/partimag/btzone/myimg/sda5/\", assign the IP address 192.168.120.183 for the tracker's URL, run"
    echo "  $ocs -p 192.168.120.183 start myimg sda1 sda5"
    echo "To stop the bittorrent service for the image dir \"/home/partimag/btzone/myimg/sda5/\", run"
    echo "  $ocs stop my sda5"
    echo
} # end of USAGE
#
get_available_bt_port() {
  while nc -w 1 -z localhost $port_b; do
    # It's used. Use the next one by adding 2.
    port_b="$((port_b+2))"
  done
} # end of get_available_bt_port
#
wait_bt_port_open() {
  local port_query="$1"
  echo -n "Waiting for port $port_query of tracker to be started successfully..."
  while ! nc -w 1 -z localhost $port_query; do
    sleep 0.5
  done
  echo " done!"
} # end of wait_bt_port_open
#
get_idir_name() {
   local t_dir="$1"
   # Return global variable "idir_name"
   # t_dir is like: xenial-x64-20161104/sda1 xenial-x64-20161104/sda5
   t_dir_imgname="$(dirname $t_dir)"
   t_dir_devname="$(basename $t_dir)"
   if [ "$t_dir_imgname" = "." ]; then
     t_dir_imgname=""
   fi
   idir_name="${t_dir_imgname}~${t_dir_devname}"  # make it like: xenial-x64-20161104~sda1
} # end of get_idir_name
#
start_bt_srv() {
  local ipt idir idir_name ezio_log_idir
  local cpu_no cmd_make_metainfo cmd_update_metainfo
  local existing_ip existing_port 
  local cmd_start_seeding ezio_start_seeding_param ezio_ram_size_KB ct_pid_log_prefix
  if [ -z "$img_dir" ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "No image dir was assigned!"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    echo "$msg_program_stop!"
    exit 1
  fi
  if [ -z "$parts" ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "Neither partition nor LV was assigned! No way to start BT service."
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    echo "$msg_program_stop!"
    exit 1
  fi
  if [ "$bt_slices_from_dev" = "img" ]; then
    # Check if the bt image dir exists
    for ipt in $parts; do
      if [ ! -d "$ocsroot_btzone/$img_dir/$(to_filename $ipt)" ]; then
        [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
        echo "Image dir \"$ocsroot_btzone/$img_dir/$(to_filename $ipt)\" not found!"
        [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
        echo "$msg_program_stop!"
        exit 1
      fi
    done
  fi
  # Stop previous bt service first.
  stop_bt_srv

  # Assign an initial port no.
  port_b="${BT_PORT_INIT}"

  if [ -z "$btsrv_ip" ]; then
    # Decide the BT server IP address if it's not assigned.
    btsrv_ip="$(get-all-nic-ip -b | awk -F" " '{print $1}')"
  fi
  echo "Using IP address \"$btsrv_ip\" for BT server..."

  ezio_start_seeding_param=""
  ct_pid_log_prefix=""
  for ipt in $parts; do
    # Reset cmd_start_seeding for each ipt, because for ezio,
    # we will collect all param and run it outside the for loop.
    cmd_start_seeding=""
    # Prepare port no
    get_available_bt_port # Obtain port_b

    idir="$img_dir/$(to_filename $ipt)" # e.g., xenial-x64-20161104/sda1
    # idir_name is the name for log usage
    get_idir_name $idir # Obtian $idir_name, e.g., xenial-x64-20161104~sda1
    # Part 1: create metainfo file (.torrent)
    echo $msg_delimiter_star_line
    # Set the variable cmd_make_metainfo
    if [ "$metainfo_creator" = "ctorrent" ]; then
      cmd_make_metainfo="ctorrent -t -p -c For_Clonezilla -l ${p_length_ctorrent} -u http://${btsrv_ip}:$port_b/announce -s $ocsroot_btzone/${idir}.torrent $ocsroot_btzone/$idir"
    elif [ "$metainfo_creator" = "transmission-create" ]; then
      cmd_make_metainfo="transmission-create -s ${p_length_transmission} -p -c For_Clonezilla -t http://${btsrv_ip}:$port_b/announce -o $ocsroot_btzone/${idir}.torrent $ocsroot_btzone/$idir"
    else
      cpu_no="$(LC_ALL=C grep -Ew "^processor[[:space:]]+" /proc/cpuinfo | wc -l)"
      cmd_make_metainfo="mktorrent -p -c For_Clonezilla -t ${cpu_no} -l ${p_length_mktorrent} -a http://${btsrv_ip}:$port_b/announce -o $ocsroot_btzone/${idir}.torrent $ocsroot_btzone/$idir"
    fi
    if [ ! -e $ocsroot_btzone/${idir}.torrent ]; then
      echo "Generate metainfo file $ocsroot_btzone/${idir}.torrent by:"
      echo "$cmd_make_metainfo"
      eval $cmd_make_metainfo
    else
      # Maybe server IP address is changed. If so, regenerate it.
      # E.g. ...:announce35:http://192.168.22.250:6969/announce7:...
      existing_ip="$(LC_ALL=C strings $ocsroot_btzone/${idir}.torrent | grep -Eo ":announce[[:digit:]]+:http://[[:digit:]\.:]+.*announce[[:digit:]]+:" | awk -F":" '{print $4}' | sed -r -e "s|\/\/||g")"
      existing_port="$(LC_ALL=C strings $ocsroot_btzone/${idir}.torrent | grep -Eo ":announce[[:digit:]]+:http://[[:digit:]\.:]+.*announce[[:digit:]]+:" | awk -F":" '{print $5}' | awk -F"/" '{print $1}')"
      if [ "${existing_ip}" != "${btsrv_ip}" ]; then
        [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
        #echo "Regenerating metainfo due to existing tracker \"${existing_ip}\" in $ocsroot_btzone/${idir}.torrent is different from detected one:\"${btsrv_ip}\"..."
        #[ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
	#rm -f $ocsroot_btzone/${idir}.torrent
        #echo "$cmd_make_metainfo"
        #eval $cmd_make_metainfo

	# Updating URL using transmission-edit
	# d8:announce34:http://192.168.7.254:6969/announce7:comment14:For_Clonezilla10:... ->
	# d8:announce34:http://10.0.1.254:6969/announce7:comment14:For_Clonezilla10:...
        [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
        echo "Updating metainfo due to existing tracker \"${existing_ip}\" in $ocsroot_btzone/${idir}.torrent is different from detected one:\"${btsrv_ip}\"..."
        [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
	cmd_update_metainfo="transmission-edit -r http://${existing_ip}:${existing_port}/announce http://${btsrv_ip}:$port_b/announce $ocsroot_btzone/${idir}.torrent"
	echo "$cmd_update_metainfo"
	eval $cmd_update_metainfo
      fi
    fi

    # Part 2: start tracker
    echo $msg_delimiter_star_line
    d_file="$btlog_dir/${idir_name}-dfile"
    log_file="$btlog_dir/${idir_name}-tracker.log"
    tracker_pid_log="$btlog_dir/${idir_name}-tracker.pid"
    if [ -n "$btclient_no" ]; then
      btclient_no_opt="--tracker_max_completed $btclient_no"
    fi
    rm -f $d_file $log_file $tracker_pid_log
    # cmd_start_tracker="ocs-bttrack --port $port_b --dfile $d_file $btclient_no_opt --logfile $log_file --nat_check 0 --scrape_allowed full --allowed_dir $(dirname $(readlink -f $ocsroot_btzone/${idir}.torrent)) &"
    cmd_start_tracker="ocs-bttrack --port $port_b --dfile $d_file $btclient_no_opt --logfile $log_file --nat_check 0 --scrape_allowed full &"
    # Squeeze multiple spaces to one only
    cmd_start_tracker="$(echo "$cmd_start_tracker" | sed -r -e "s|[[:space:]]+| |g")"
    echo "Start tracker by:"
    echo "$cmd_start_tracker"
    eval $cmd_start_tracker
    tracker_pid="$!"
    echo "$tracker_pid" > $tracker_pid_log
    wait_bt_port_open $port_b

    # Part 3: start bitorrent client program
    echo $msg_delimiter_star_line
    ct_pid_log="$btlog_dir/${idir_name}-torrent.pid" # Used for non-ezio bt client
    if [ "$bt_client" = "lftp" ]; then
      cmd_start_seeding="lftp -c torrent -O $(dirname $ocsroot_btzone/${idir}) $ocsroot_btzone/${idir}.torrent"
    elif [ "$bt_client" = "ezio" ]; then
      # Make the pid file for ezio like: sda1~sda2~sda3-torrent.pid
      ct_pid_log_prefix="${ct_pid_log_prefix}~$(to_filename $ipt)"
      case "$bt_slices_from_dev" in
      img)
        # Collect multiple torrent files. Make it like:
	# ezio -U -T sda1.torrent -L /home/partimag/btzone/myimg/sda1/ -T sda2.torrent -L /home/partimag/btzone/myimg/sda2/
        ezio_start_seeding_param="$ezio_start_seeding_param -T $ocsroot_btzone/${idir}.torrent -L $ocsroot_btzone/$(dirname ${idir})"
        ;;
      raw-dev)
        # Squeeze multiple spaces to one only
        # Collect multiple torrent files. Make it like:
	# ezio -U -T sda1.torrent -L /dev/sda1 -T sda2.torrent -L /dev/sda2
        ezio_start_seeding_param="$ezio_start_seeding_param -T $ocsroot_btzone/${idir}.torrent -L /dev/$ipt"
        ;;
      esac
    elif [ "$bt_client" = "btdownloadheadless" ]; then
      cmd_start_seeding="btdownloadheadless --saveas $ocsroot_btzone/${idir} $ocsroot_btzone/${idir}.torrent &"
    elif [ "$bt_client" = "aria2c" ]; then
      # //NOTE// The working dir is the parent dir.
      cmd_start_seeding="aria2c -D --seed-ratio=0.0 --enable-dht6=false --enable-dht=false --bt-seed-unverified $aria2c_extra_opt -d $ocsroot_btzone/$(dirname ${idir}) $ocsroot_btzone/${idir}.torrent"
    else
      # Check the option -M of ctorrent, if it only allows 20-1000, modify it as "-M 20" from "-M 4", for example.
      # DRBL patched ctorrent allows -M to be 4-1000, otherwise it must be within 20-1000.
      # //NOTE// Here the white space after echo is a must. To avoid something like: echo "-n" -> make it as: echo " -n".
      ct_M="$(echo " $ctorrent_extra_opt" | grep -Ewo -- "[^[:space:]]*-M[[:space:]]+[[:digit:]]+[[:space:]]+" | awk -F" " '{print $2}')"
      if [ -n "$ct_M" ]; then
        if [ -n "$(LC_ALL=C strings `which ctorrent` | grep -Ew -- "^-%c argument must be between 20 and 1000")" ]; then
          # Lower limit for -M option is 20.
          if [ "$ct_M" -lt 20 ]; then
            # Force to make it as 20
            [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
            echo "The argument assigned in ctorrent_extra_opt for -M is $ct_M, which is lower than 20, the value allowed in the version of original ctorrent. Force to make it as 20."
            [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
	    ctorrent_extra_opt="$(echo " $ctorrent_extra_opt" | sed -r -e "s|-M[[:space:]]+[[:digit:]]+[[:space:]]+|-M 20 |" | sed -r -e "s|^[[:space:]]*||g")"
	  fi
        fi
      fi
      cmd_start_seeding="ctorrent -f -d $ctorrent_extra_opt -s $ocsroot_btzone/${idir} $ocsroot_btzone/${idir}.torrent"
    fi
    if [ -n "$cmd_start_seeding" ]; then
      # For non-ezio cases
      echo "Start seeding by:"
      echo "$cmd_start_seeding"
      eval $cmd_start_seeding
      ct_pid="$!"
      echo "$ct_pid" > $ct_pid_log
      echo $msg_delimiter_star_line
    fi
  done
  # For ezio case
  if [ -n "$ezio_start_seeding_param" ]; then
     get_ezio_prog  # get the variable ezio_prog
     # //NOTE// The working dir for ezio is the working dir as torrent file
     # E.g., ezio -U -f /home/partimag/btzone/xenial-x64-20161104/sda1.torrrent /home/partimag/btzone/xenial-x64-20161104
     # Set memory limit
     ezio_ram_size_KB="$(LC_ALL=C vmstat -s -S k |grep "free memory" | awk -F" " '{print $1}')"
     # Do not use all the free memory. ezio_cache_ratio is from drbl-ocs.conf.
     ezio_ram_size_KB="$(LC_ALL=C printf "%.0f" "$(echo "($ezio_ram_size_KB*$ezio_cache_ratio)" | bc -l)")"
     # Log and pid
     ezio_log_idir="${ezio_seeding_log}-${ct_pid_log_prefix}.log"
     # Remove the leading ~
     ct_pid_log_prefix="$(echo $ct_pid_log_prefix | sed -e "s/^~//g")"
     ct_pid_log="$btlog_dir/${ct_pid_log_prefix}-torrent.pid"
     case "$bt_slices_from_dev" in
     img)
       # E.g., ezio -U -f -T sda1.torrent -L /dev/sda1 -T sda2.torrent -L /dev/sda2
       cmd_start_seeding="$ezio_prog $ezio_seeder_opt $ezio_common_opt -U -f --cache $ezio_ram_size_KB $ezio_start_seeding_param" 
       ;;
     raw-dev)
       # E.g., -U -T sda1.torrent -L /dev/sda1 -T sda2.torrent -L /dev/sda2
       cmd_start_seeding="$ezio_prog $ezio_seeder_opt $ezio_common_opt -U --cache $ezio_ram_size_KB $ezio_start_seeding_param" 
       ;;
     esac
     # Squeeze multiple spaces to one only
     cmd_start_seeding="$(echo $cmd_start_seeding | sed -r -e "s|[[:space:]]+| |g")"
     cmd_start_seeding="$cmd_start_seeding > $ezio_log_idir 2>&1 &"
     echo "Start seeding by:"
     echo "$cmd_start_seeding"
     eval $cmd_start_seeding
     ct_pid="$!"
     echo "$ct_pid" > $ct_pid_log
     [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
     echo "$ezio_prog seeding status output for ${idir}.torrent will be in \"$ezio_log_idir\". Use \"tail -f $ezio_log_idir\" to check the status."
     [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
     echo $msg_delimiter_star_line
   fi
} # end of start_bt_srv
#
stop_bt_srv() {
  local id_file=""
  local all_pid rc kill_cmd ipt idir_name
  local ct_pid_log_prefix
  # List the possible id_file, and this is only list for later to test if it exists.
  case "$img_dir" in 
     "all"|"") id_file="$btlog_dir/*.pid";;
            *)
               ct_pid_log_prefix=""
  	       for ipt in $parts; do
		 idir="$img_dir/$(to_filename $ipt)"
                 # idir_name is the name for log usage
                 get_idir_name $ipt # Obtian $idir_name
                 ct_pid_log_prefix="${ct_pid_log_prefix}~$(to_filename $ipt)"
		 # Possible tracker pid file and torrent pid file, there is ezio part later.
                 id_file="$btlog_dir/${idir_name}-tracker.pid $btlog_dir/${idir_name}-torrent.pid"
               done
	       # Append Ezio part
               # Remove the leading ~
               ct_pid_log_prefix="$(echo $ct_pid_log_prefix | sed -e "s/^~//g")"
               id_file="$id_file $btlog_dir/${ct_pid_log_prefix}-tracker.pid $btlog_dir/${ct_pid_log_prefix}-torrent.pid"
	       ;;
  esac

  # All
  if [ "$brutal_kill" = "no" ]; then
    echo "Stop service: ocs-bttrack and ezio/ctorrent/lftp/btdownloadheadless/aria2c if available..."
    for i in $id_file; do
      [ ! -e "$i" ] && continue
      all_pid="$all_pid $(cat $i)"
    done
    if [ -n "$all_pid" ]; then
      echo "Found pid..."
      kill_cmd="kill -9 $all_pid"
      echo "Running: $kill_cmd"
      eval $kill_cmd
      rc=$?
      if [ "$rc" -eq 0 ]; then
        rm -f $id_file
      fi
    else
      echo "No pid was found... Do nothing."
    fi
  else
    echo "Brutally kill service ocs-bttrack and ezio/ctorrent/lftp/btdownloadheadless/aria2c if available..."
    pkill -9 ocs-bttrack
    pkill -9 ezio
    pkill -9 ezio-static
    pkill -9 ctorrent
    pkill -9 lftp
    pkill -9 btdownloadheadless
    pkill -9 aria2c
  fi
  echo "Done!"
} # end of stop_bt_srv
#
show_status() {
  echo $msg_delimiter_star_line
  ps -www -C "ocs-bttrack" -C "aria2c" -C "ctorrent" -C "lftp" -C "btdownloadheadless" -o pid,tname,cmd
  echo $msg_delimiter_star_line
} # end of show_status

################
##### MAIN #####
################

ocs_file="$0"
ocs=`basename $ocs_file`
#
while [ $# -gt 0 ]; do
 case "$1" in
   -b|--brutal-kill) brutal_kill="yes"
	   shift;;
   -c|--metainfo-creator) 
           shift; 
           if [ -z "$(echo $1 |grep ^-.)" ]; then
             # skip the -xx option, in case 
	     metainfo_creator="$1"
             shift;
           fi
           [ -z "$btlog_dir" ] && USAGE && exit 1
           ;;
   -f|--from-raw-dev) bt_slices_from_dev="raw-dev"
	   shift;;
   -g|--log-dir)
           shift; 
           if [ -z "$(echo $1 |grep ^-.)" ]; then
             # skip the -xx option, in case 
             btlog_dir="$1"
             shift;
           fi
           [ -z "$btlog_dir" ] && USAGE && exit 1
           ;;
   -m|--max-client-no)
           shift; 
           if [ -z "$(echo $1 |grep ^-.)" ]; then
             # skip the -xx option, in case 
             btclient_no="$1"
             shift;
           fi
           [ -z "$btclient_no" ] && USAGE && exit 1
           ;;
   -p|--bt-srv-ip)
           shift; 
           if [ -z "$(echo $1 |grep ^-.)" ]; then
             # skip the -xx option, in case 
             btsrv_ip="$1"
             shift;
           fi
           [ -z "$btsrv_ip" ] && USAGE && exit 1
           ;;
   -r|--bt-root)
           shift; 
           if [ -z "$(echo $1 |grep ^-.)" ]; then
             # skip the -xx option, in case 
             ocsroot_btzone="$1"
             shift;
           fi
           [ -z "$ocsroot_btzone" ] && USAGE && exit 1
           ;;
   -t|--bt-client)
           shift; 
           if [ -z "$(echo $1 |grep ^-.)" ]; then
             # skip the -xx option, in case 
             bt_client="$1"
             shift;
           fi
           [ -z "$bt_client" ] && USAGE && exit 1
           ;;
   -*)     echo "${0}: ${1}: invalid option" >&2
           USAGE >& 2
           exit 2 ;;
   *)      break ;;
 esac
done

action="$1"
shift
img_dir="$1"
shift
parts="$*"

# Load deafult values if not specified.
[ -z "$metainfo_creator" ] && metainfo_creator="$metainfo_creator_def"
[ -z "$bt_client" ] && bt_client="$bt_client_def"
if [ "$bt_slices_from_dev" = "raw-dev" ]; then
  [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
  echo "Option -f or --from-raw-dev is assigned. Force to use ezio as the BT client."
  [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
  bt_client="ezio"
fi

# Strip the leading /dev/ if it's assigned. Make it like sda1, sdb1, not /dev/sda1, /dev/sdb1.
parts="$(echo $parts | sed -r -e "s|/dev/||g")"  # partition name, e.g. sda1, sda2

#check_if_root
ask_and_load_lang_set

case "$action" in
   start)  start_bt_srv;;
    stop)  stop_bt_srv;;
  status)  show_status;;
       *)  USAGE; exit 1;;
esac
