Bash script to swap out, edit host files

Question 1

This is my first useful bash script. It’s meant to make it easy to switch between a “work” hosts file (that blocks sites like reddit.com, youtube.com, etc.) and my normal hosts file, and also to allow me to easily add/remove sites I want to block.

I’m not sure if I’m doing everything right or securely, so please feel free to pick at anything.

#!/bin/bash

usage()
{
cat <<EOF
Usage: wrk [command] [host1 host2 ...]

Commands:
    list                list blocked hosts
    add [host ...]      add a host to be blocked
    rm [host ...]       remove hosts from block
    start               start blocking
    stop                stop blocking
EOF
}

VARDIR="$HOME/abes_commands/var"    
# HOST_FILE - actual host file to be swapped around.
# ORIG_FILE - original host file without any of the appended blocked hosts.
# BLCK_FILE - list of blocked hosts.
# BLCK_TEMP - temporary file used when removing hosts.
HOST_FILE="/etc/hosts"
ORIG_FILE="$VARDIR/original_host"
BLCK_FILE="$VARDIR/blocked_host"
BLCK_TEMP=`mktemp -t "blocked_hosts"` || `mktemp /tmp/blocked_hosts.XXXXXXX` || exit 1

# Make sure files exist.
[[ -e $ORIG_FILE ]] || touch "$ORIG_FILE"
[[ -e $BLCK_FILE ]] || touch "$BLCK_FILE"

# Check to see if the block is currently active.
ACTIVE_FLAG="$HOME/.wrk_block.flag"
if [[ -e $ACTIVE_FLAG ]]; then 
    IS_ACTIVE=0
else 
    IS_ACTIVE=1
fi

add_host()
{
    local hosts=("$@")
    for host in "${hosts[@]:1}"; do
        # append host to blocked hosts list
        echo "127.0.0.1 $host" >> "$BLCK_FILE"
        echo -e "33[0;32madded33[0m $host" 
    done
}

# todo: remove need for loop; do in one go
rm_host()
{
    local hosts=("$@")
    for host in "${hosts[@]:1}"; do
        # overwrite host list file with a copy removing a certain host
        awk -v host=$host 'NF==2 && $2!=host { print }' "$BLCK_FILE" > "$BLCK_TEMP"
        mv "$BLCK_TEMP" "$BLCK_FILE"
        echo -e "33[0;31mremoved33[0m $host" 
    done
}

check_root()
{
    if [[ "`whoami`" != "root" ]]; then
        echo "You don't have sufficient priviledges to run this script (try sudo.)"
        exit 1
    else 
        [[ -e $HOST_FILE ]] || { echo "Can't find or access host file."; exit 1; }
    fi
}

start_block()
{
    if [[ $IS_ACTIVE -ne 0 ]]; then
        cp "$HOST_FILE" "$ORIG_FILE"
        cat "$BLCK_FILE" >> "$HOST_FILE"
        touch "$ACTIVE_FLAG"
        echo "Block started."
    else
        echo "Already blocking."
    fi
}

stop_block()
{
    if [[ $IS_ACTIVE -eq 0 ]]; then
        cp "$ORIG_FILE" "$HOST_FILE"
        [[ -e $ACTIVE_FLAG ]] && rm "$ACTIVE_FLAG"
        echo "Stopped blocking."
    else
        echo "Not blocking."
    fi
}

case $1 in
    'ls' | 'list') 
        awk 'NF == 2 { print $2 }; END { if (!NR) print "Empty" }' "$BLCK_FILE";;
    'add') 
        [[ -z $2 ]] && { usage; exit 1; }
        add_host $@;;
    'rm' | 'remove') 
        [[ -z $2 ]] && { usage; exit 1; }
        rm_host $@;;
    'start') 
        check_root
        start_block;;
    'stop') 
        check_root
        stop_block;;
    *)
        usage;;
esac

The blocked_host file is appended to the hosts file when blocking starts, and looks like this:

127.0.0.1 www.reddit.com
127.0.0.1 reddit.com
127.0.0.1 www.news.ycombinator.com
127.0.0.1 news.ycombinator.com

Question 2

Your code is written very clearly and easy to read.

I am only proficient in ksh programming but as far as I can see in your subroutines you loop through your argument list skipping the first entry. if you augment your command syntax and the hostlist now starts with the third argument you have to change the script at a lot of different places. this is annoying and error prone. so I would prefer assigning the arguments to named parameters as early as possible at least for all but the hosts, so
```
COMMAND=$1
shift
```
I would write data to stdout but messages to stderr. Especially error messages or warnings should be written to stderr. the usage message also should be written to standard error

I doubt that

BLCK_TEMP=`mktemp -t "blocked_hosts"` || `mktemp /tmp/blocked_hosts.XXXXXXX` || exit 1

will work as you expected. You should test it.

why not always use the directory mktemp /tmp/blocked_hosts.XXXXXXX or better mktemp -t blocked_hosts.XXXXXXX.
if mktemp -t "blocked_hosts" is alreday in use why do you want to use mktemp /tmp/blocked_hosts.XXXXXXX (maybe this is in a different directory) and not mktemp -t blocked_hosts.XXXXXXX?

I would write all constants a block of contiguous lines so

HOST_FILE="/etc/hosts"
ORIG_FILE="$VARDIR/original_host" 
ACTIVE_FLAG="$HOME/.wrk_block.flag"

The they are easy to find and to modify

I would prefer error-codes different from 1 but errorcodes specific for my application. 1 is often generated by other commands. Also I would use different error codes for different errors.
If I am using a command I did not need any messages that the command was successfully executed. Tools like cron generate mails if a command produces messages and mails them to the user. messages are necessary if something unexpected happens that perhaps needs special actions from the user. if the command executes successfully and nothing special happens this should not be told to the user or the calling program.
I don’t like/need/want coloured messages. Often I am not able to read them.
If something unexpected happens I try to exit the script. This can be achieved by setting -e or a trap for ERR. In this case I remove the created files too.
Instead of
```
Usage: wrk [command] [host1 host2 ...] 
```
you can do
```
Usage: $(basename $0) [command] [host1 host2 ...] 
```
If someone renames the file the appropriate file name is substituted in the usage message. But I am not sure if one should support the renaming of the file.
in the start_blocking function:
```
cp "$HOST_FILE" "$ORIG_FILE"
cat "$BLCK_FILE" >> "$HOST_FILE"
touch "$ACTIVE_FLAG" 
```
I would first do the lock (create the $ACIVE_FLAG file) and then do my actions.
in the stop_blocking function:
```
   [[ -e $ACTIVE_FLAG ]] && rm "$ACTIVE_FLAG" 
```
if there is no $ACTIVE_FLAG I think that is worth a message to the user. I am not sure If you should proceed in this case and copy to the hosts file
Related to:
```
cp "$HOST_FILE" "$ORIG_FILE" 
cat "$BLCK_FILE" >> "$HOST_FILE" 
```
why cpin the first line and cat with >> in the next line. Is there a difference?
will $ACTIVE_FLAG always be the same, also if a user executes the scrips with sudo or su(i don’t know much about this). I think it should be a directory independent from the user and user settings
you should test your code so that each command is executed at least once. your code can be tested easily because you can change the global variables so that you do no actually change a /etc/hosts file

Question 3

I think this is generally clean.

Style Points

use $() instead of ``, it’s clearer and they can be nested.

E.g:

if [[ "$(whoami)" != "root" ]]; then
    echo "You don't have sufficient priviledges to run this script (try sudo.)"
    exit 1

Consider setting:

set -u #prevent unset variables
set -e #stop on an error

At the top of your script.

rm_host may benefit from rewriting using an inplace sed edit, rather than using awk to copy over the top of the old file.

Locks

Although unlikely, you may want to consider a lock to prevent races between add_host and rm_host. $ACTIVE_FLAG would make a good lock.

$VARDIR

It looks like you intend this script to be used at differing privileges level (e.g. root and non-root). This will change the $HOME env variable and consequently the $VARDIR env variable. Potentially {add,rm}_host may act on different files to {start,stop}_block. Note that this won’t present itself if you only use sudo.

Reconsider what you want $VARDIR to be. On a a single user machine /var/ may be a good choice.

Question 4

For ls, do not print Empty if there is no hosts blocked. The correct way is to be silent if you find no data. This will (assuming some one else uses it) allow them not to special case the Empty case. That is the below should suffice.
case ls | list)
cut -f2 -d < “$BLCK_FILE”;;

Another style point. You might want to consider avoiding the superfluous quotes. That is, use case ls) rather than case ‘ls’). I think the former is more often used, and clearer.

Consider also that it could be used from a cron (to switch at particular times of the day.) So it might be better to assume that you don’t have access to $HOME (As mentioned in the above comment)

{start,stop}_block consider using exit 1 to indicate error. Perhaps

start_block()
{
    [[ $IS_ACTIVE -ne 0 ]] || {exit 1}
    cp "$HOST_FILE" "$ORIG_FILE"
    cat "$BLCK_FILE" >> "$HOST_FILE"
    touch "$ACTIVE_FLAG"
    echo "Block started."
}

You might also want to think what the behavior would be, when you add a new host when a block is already in effect.

I am not sure why you would need this. You are already doing this in start_block

# Make sure files exist.
[[ -e $ORIG_FILE ]] || touch "$ORIG_FILE"

It is cleaner to use id -g rather than whoami.

Your check_root may be better expressed as below because this is the only check you need to do.

check_root()
{
    if [[ -w $HOST_FILE ]]; then
        echo "You don't have sufficient priviledges to access $HOST_FILE"
        exit 1
    fi
}

It is considered a better practice in Unix to be silent when there is not much information to be added.

# explode args into lines.
argx() {
    for var in "$@"; do
       echo $var
    done
}
add_host()
{
    argx | sed -e 's/^/127.0.0.1 /g' >> "$BLCK_FILE"
}

# you can use comm to print a list of removed files.
rm_host()
{
    cat "$BLCK_FILE" |sort -u |comm -13 <(argx |sort -u) - > "$BLCK_TEMP"
    mv "$BLCK_TEMP" "$BLCK_FILE"
}

You can reduce active flag checking to

ACTIVE_FLAG="$HOME/.wrk_block.flag"
IS_ACTIVE=$(test -e $ACTIVE_FLAG; echo $?)

Style Points

Locks

$VARDIR

Leave a Reply Cancel reply