gget/gget.c

/*
  Copyright 2014 Grégory Soutadé

  This file is part of gget.

  gget is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  gget is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with gget.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <sys/types.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <fcntl.h>
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <curl/curl.h>
#include <libgen.h>

#define DEFAULT_USER_AGENT "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:30.0) Gecko/20100101 Firefox/30.0"
#define DEFAULT_NB_THREADS 3
#define MAX_NB_THREADS 10
#define DEFAULT_OUT_FILENAME "gget.out"

#ifdef WIN32
#include <windows.h>
#define GNU_ERR ""

typedef HANDLE pthread_t;
typedef void pthread_attr_t;

static int get_console_width()
{
    RECT r;
    HWND console = GetConsoleWindow();

    GetWindowRect(console, &r);

    return r.right-r.left;
}

static int pthread_create(pthread_t *thread, const pthread_attr_t *attr,
                          void *(*start_routine) (void *), void *arg)
{
    *thread = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)start_routine,
                           arg, 0, NULL);

    return 0;
}

static int pthread_join(pthread_t thread, void **retval)
{
    WaitForSingleObject(thread, 0);
    GetExitCodeThread(thread,(LPDWORD)retval);
    CloseHandle(thread);

    return 0;
}

#else
#include <sys/ioctl.h>
#include <pthread.h>

#define GNU_ERR " (%m)"
static int get_console_width()
{
    struct winsize w;

    ioctl(STDOUT_FILENO, TIOCGWINSZ, &w);

    return w.ws_col;
}
#endif

typedef struct {
    curl_off_t dltotal;
    curl_off_t dlnow;
    unsigned speed;
} stats_t ;

typedef struct {
    pthread_t thread;
    char* tmp_filename;
    //
    CURL *curl;
    char* url;
    char* user_agent;
    int fd;
    int id;
    unsigned already_downloaded;
    unsigned start;
    unsigned end;
    curl_off_t max_speed;
    stats_t* stats;
} transfert_t;

typedef struct {
    unsigned nb_threads;
    stats_t* stats;
    unsigned end;
} stats_params_t;

static void* display_stats(stats_params_t* p)
{
    unsigned speed, percent, time_left, max_time_left;
    unsigned total_percent, hours, minutes, seconds;
    char* suffix;
    stats_t* cur;
    int i, nb_chars_printed;

    while (!p->end)
    {
        // If the window has been resized
        max_time_left = 0;
        suffix = "B";
        total_percent = 0;
        nb_chars_printed = 0;

        printf("\r");
        for(i=0; i<p->nb_threads; i++)
        {
            cur = &p->stats[i];

            speed = cur->speed;

            if (cur->dltotal)
                percent = (cur->dlnow*100)/cur->dltotal;
            else
                percent = 0;

            if (speed)
                time_left = (cur->dltotal-cur->dlnow)/speed;
            else
                time_left = 0;

            if (speed < (1024*1024))
            {
                speed /= 1024;
                suffix = "kB";
            }
            else
            {
                speed /= 1024*1024;
                suffix = "MB";
            }

            nb_chars_printed += 
                printf("T%d: %02u%% %u%s/s  ",
                       i, percent, speed, suffix);

            if (time_left > max_time_left)
                max_time_left = time_left;

            total_percent += percent/p->nb_threads;
        }

        nb_chars_printed += printf("  Total: %u%% ", total_percent);

        if (max_time_left < 60)
        {
            nb_chars_printed += printf("eta %us", max_time_left);
        }
        else
        {
            seconds = max_time_left % 60;
            max_time_left /= 60;

            if (max_time_left < 60)
            {
                minutes = max_time_left;
                nb_chars_printed += printf("eta %um %us", minutes, seconds);
            }
            else
            {
                minutes = max_time_left % 60;
                hours = max_time_left / 60;
                nb_chars_printed += printf("eta %uh %um %us", hours, minutes, seconds);
            }
        }

        for(i=nb_chars_printed; i<get_console_width(); i++)
            printf(" ");

        fflush(stdout);

        sleep(1);
    }

    return NULL;
}

static size_t write_cb(void *contents, size_t size, size_t nmemb, void *userp)
{
    int ret;
    ret = write(((transfert_t*)userp)->fd, contents, size*nmemb);
    if (ret < 0)
        printf("Error write" GNU_ERR "\n");
    return ret;
}

static int progress_cb(void *clientp, curl_off_t dltotal, curl_off_t dlnow, curl_off_t ultotal, curl_off_t ulnow)
{
    transfert_t* t = (transfert_t*)clientp;
    double speed_d;

    curl_easy_getinfo(t->curl, CURLINFO_SPEED_DOWNLOAD, &speed_d);

    t->stats[t->id].dltotal = dltotal + t->already_downloaded;
    t->stats[t->id].dlnow = dlnow + t->already_downloaded;
    t->stats[t->id].speed = speed_d;

    return 0;
}

void* do_transfert(transfert_t* t)
{
    CURLcode res;
    char range[50];

    snprintf(range, sizeof(range), "%u-%u", t->start, t->end);

    curl_easy_setopt(t->curl, CURLOPT_USERAGENT, t->user_agent);
    curl_easy_setopt(t->curl, CURLOPT_URL, t->url);
    curl_easy_setopt(t->curl, CURLOPT_FOLLOWLOCATION, 1L);
    curl_easy_setopt(t->curl, CURLOPT_RANGE, range);

    if (t->max_speed)
        curl_easy_setopt(t->curl, CURLOPT_MAX_RECV_SPEED_LARGE, t->max_speed);

    /* send all data to this function  */
    curl_easy_setopt(t->curl, CURLOPT_WRITEFUNCTION, write_cb);
 
    /* we pass our 'chunk' struct to the callback function */
    curl_easy_setopt(t->curl, CURLOPT_WRITEDATA, (void*)t);

    curl_easy_setopt(t->curl, CURLOPT_NOPROGRESS, 0L);
    curl_easy_setopt(t->curl, CURLOPT_XFERINFOFUNCTION, progress_cb);
    curl_easy_setopt(t->curl, CURLOPT_XFERINFODATA, t);

    /* Perform the request, res will get the return code */ 
    res = curl_easy_perform(t->curl);
    /* Check for errors */ 
    if(res != CURLE_OK)
        fprintf(stderr, "curl_easy_perform() failed: %s\n",
                curl_easy_strerror(res));

    return NULL;
}

static int configure_transfert(int id, transfert_t* t,
                               char* url, char* filename,
                               unsigned start, unsigned end,
                               unsigned max_speed, char* user_agent,
                               int* exists)
{
    // filename + . + number + \0
    unsigned filename_size = strlen(filename)+1+3+1;
    struct stat s;

    t->curl = curl_easy_init();
    
    if(!t->curl)
        return -1;

    t->url = url;

    t->tmp_filename = malloc(filename_size);
    snprintf(t->tmp_filename, filename_size, "%s.%d",
             filename, id);
    t->id = id;

    t->start = start;
    t->end = end;
    t->max_speed = max_speed;
    t->user_agent = user_agent;

    if (stat(t->tmp_filename, &s))
    {
        *exists = 0;

        t->fd = open(t->tmp_filename,
                     O_WRONLY | O_CREAT,
                     S_IRUSR|S_IWUSR);
        t->already_downloaded = 0;
    }
    else
    {
        *exists = 1;

        t->fd = open(t->tmp_filename,
                     O_WRONLY | O_APPEND);

        t->start += s.st_size;
        t->already_downloaded = s.st_size;
    }

    if (!t->fd)
    {
        printf("Opening '%s' failed" GNU_ERR "\n",
               t->tmp_filename);
        return -1;
    }
    
    return 0;
}

static int free_transfert(transfert_t* t)
{
    if (t)
    {
        if (t->curl)
            curl_easy_cleanup(t->curl);
        
        if (t->tmp_filename)
            free(t->tmp_filename);
    }

    return 0;
}

static void merge_files(transfert_t* transferts, int nb_threads,
                        char* out_filename)
{
    int fd, fd_tmp, i;
    struct stat s;
    char* buffer;

    rename(transferts[0].tmp_filename, out_filename);
    fd = open(out_filename, O_WRONLY | O_APPEND);

    for(i=1; i<nb_threads; i++)
    {
        fd_tmp = open(transferts[i].tmp_filename, O_RDONLY);
        stat(transferts[i].tmp_filename, &s);
        buffer = malloc(s.st_size);
        read(fd_tmp, buffer, s.st_size);
        write(fd, buffer, s.st_size);
        free(buffer);
        close(fd_tmp);
        unlink(transferts[i].tmp_filename);
    }
    close(fd);
}

static int get_file_info(char* url, char* user_agent,
                         char** filename, unsigned* size,
                         unsigned quiet)
{
    double length;
    CURLcode res;
    int ret = 0;
    unsigned response_code;
    char* real_name = NULL;
    CURL* curl = curl_easy_init();

    if(!curl)
        return -1;

    curl_easy_setopt(curl, CURLOPT_USERAGENT, user_agent);
    curl_easy_setopt(curl, CURLOPT_URL, url);
    curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
    curl_easy_setopt(curl, CURLOPT_NOBODY, 1L);
    if (!quiet)
        curl_easy_setopt(curl, CURLOPT_HEADER, 1L);

    res = curl_easy_perform(curl);

    /* Check for errors */
    if(res != CURLE_OK)
    {
        fprintf(stderr, "curl_easy_perform() failed: %s\n",
                curl_easy_strerror(res));
        ret = -1;
        goto end;
    }

    curl_easy_getinfo( curl, CURLINFO_RESPONSE_CODE, &response_code);

    if (response_code != 200)
    {
        ret = -1;
        goto end;
    }

    curl_easy_getinfo( curl,
                       CURLINFO_CONTENT_LENGTH_DOWNLOAD, &length );
    *size = length;

    if (*filename == NULL)
    {
        curl_easy_getinfo( curl, CURLINFO_EFFECTIVE_URL, &real_name );

        if (real_name)
        {
            *filename = strdup(basename(real_name));

            if (!strcmp(*filename, "."))
            {
                free(*filename);
                *filename = strdup(DEFAULT_OUT_FILENAME);
                printf("Filename not found, output to %s\n", DEFAULT_OUT_FILENAME);
            }
        }
    }

end:
    curl_easy_cleanup(curl);

    return ret;
}

static void find_free_file(char** filename)
{
    struct stat s;
    int cur_index;
    char* new_filename;
    unsigned new_size;

    if (stat(*filename, &s)) return ;

    new_size = strlen(*filename) + 1 + 2 + 1;
    new_filename = malloc(new_size);

    for (cur_index = 1; cur_index < 100; cur_index++)
    {
        snprintf(new_filename, new_size, "%s.%d",
                 *filename, cur_index);

        if (stat(new_filename, &s))
        {
            free(*filename);
            *filename = new_filename;
            return ;
        }
    }

    free(*filename);
    *filename = NULL;
}

static void usage(char* program_name)
{
    printf("%s: Parallel HTTP file download\n", program_name);
    printf("usage: %s [-n nb_threads] [-l speed_limit] [-o out_filename] [-u user_agent] [-q] [-h] url\n",
        program_name);
    printf("\t-n : Specify number of threads (default : %d)\n", DEFAULT_NB_THREADS);
    printf("\t-l : Download speed limit for all threads (not per thread)\n");
    printf("\t-o : Out filename, default is retrieved by GET request or '%s' if not found\n",
           DEFAULT_OUT_FILENAME);
    printf("\t-u : User agent, default is '%s'\n", DEFAULT_USER_AGENT);
    printf("\t-q : Quiet mode\n");
    printf("\t-h : Display help\n");
}

int main(int argc, char** argv)
{
    pthread_t display_thread = 0;
    unsigned nb_threads = DEFAULT_NB_THREADS;
    char* user_agent = strdup(DEFAULT_USER_AGENT);
    stats_t* stats = NULL;
    int ret = -1, i;
    stats_params_t stats_params;
    transfert_t* transferts = NULL;
    unsigned total_size, thread_size;
    unsigned start, end, max_speed = 0, quiet = 0;
    char* out_filename = NULL, *url = NULL;
    char* suffix = "B";
    void* res;
    int opt;
    struct timeval time_start, time_end;
    struct tm * full_time;
    double average_speed;
    int continuous_mode=0;
    int exists;

    while ((opt = getopt(argc, argv, "hn:l:o:u:q")) != -1) {
        switch (opt) {
        case 'n':
            nb_threads = atoi(optarg);
            if (nb_threads == 0)
                nb_threads = DEFAULT_NB_THREADS;
            else if (nb_threads > MAX_NB_THREADS)
            {
                printf("Max numberb of threads is %d\n", MAX_NB_THREADS);
                nb_threads = MAX_NB_THREADS;
            }
            break;
        case 'l':
            max_speed = atoi(optarg);
            break;
        case 'o':
            out_filename = strdup(optarg);
            break;
        case 'u':
            user_agent = strdup(optarg);
            break;
        case 'q':
            quiet = 1;
            break;
        default:
        case 'h':
            usage(argv[0]);
            return 0;
        }
    }

    if (optind != argc-1)
    {
        usage(argv[0]);
        return 0;
    }
    else
        url = argv[optind];

    if (get_file_info(url, user_agent,
                      &out_filename, &total_size,
                      quiet))
        return -1;

    if (!out_filename)
        out_filename = strdup(DEFAULT_OUT_FILENAME);

    find_free_file(&out_filename);

    if (out_filename == NULL)
    {
        printf("Unable to find free file to write to !\n");
        goto end;
    }
    else
        printf("Save in '%s'\n\n", out_filename);

    stats = malloc(sizeof(*stats)*nb_threads);
    memset(stats, 0, sizeof(*stats)*nb_threads);

    transferts = malloc(sizeof(*transferts)*(nb_threads+1));
    memset(transferts, 0, sizeof(*transferts)*(nb_threads+1));

    thread_size = total_size/nb_threads;
    max_speed /= nb_threads;

    for(i=0; i<nb_threads; i++)
    {
        transferts[i].stats = stats;

        start = thread_size*i;
        if (i < (nb_threads-1))
            end = thread_size*(i+1)-1;
        else
            end = total_size;

        ret = configure_transfert(i, &transferts[i], url, out_filename,
                                  start, end, max_speed, user_agent, &exists);
        if (ret)
            goto end;

        // First set continuous mode
        if (i == 0)
        {
            continuous_mode = exists;
        }
        else
        {
            // Check valid continuous mode or not
            if (exists ^ continuous_mode)
            {
                printf("Error : you already started to download this file with a different number of thread, please clear temporary files or restart with the same number of threads\n");
                goto end;
            }
        }
    }

    // Check for last temporary file
    configure_transfert(i, &transferts[i], url, out_filename,
                        start, end, max_speed, user_agent, &exists);
    unlink(transferts[i].tmp_filename);

    if (exists)
    {
        printf("Error : you already started to download this file with a different number of thread, please clear temporary files or restart with the same number of threads\n");
        free_transfert(&transferts[i]);
        goto end;
    }

    stats_params.end = 0;
    stats_params.nb_threads = nb_threads;
    stats_params.stats = stats;

    if (!quiet)
    {
        pthread_create(&display_thread, NULL,
                       (void*(*)(void*))display_stats,
                       (void*) &stats_params);
    }

    gettimeofday(&time_start, NULL);

    for(i=0; i<nb_threads; i++)
    {
        pthread_create(&transferts[i].thread, NULL,
                       (void*(*)(void*))do_transfert,
                       (void*) &transferts[i]);
    }

    for(i=0; i<nb_threads; i++)
    {
        pthread_join(transferts[i].thread, &res);
    }

    gettimeofday(&time_end, NULL);

    merge_files(transferts, nb_threads, out_filename);

    if (!quiet)
    {
        full_time = localtime((const time_t*)&time_end.tv_sec);

        average_speed = (double)total_size / (double)(time_end.tv_sec - time_start.tv_sec);

        if (average_speed < (1024*1024))
        {
            average_speed /= 1024;
            suffix = "kB";
        }
        else
        {
            average_speed /= 1024*1024;
            suffix = "MB";
        }

        printf("\n\n%04d-%02d-%02d %02d:%02d:%02d (%.2f%s/s) '%s' saved\n",
               full_time->tm_year+1900, full_time->tm_mon, full_time->tm_mday,
               full_time->tm_hour, full_time->tm_min, full_time->tm_sec,
               average_speed, suffix, out_filename);
    }

    ret = 0;

end:
    if (stats)
        free(stats);

    free(user_agent);
    free(out_filename);

    stats_params.end = 1;

    for(i=0; i<nb_threads; i++)
        free_transfert(&transferts[i]);

    if (!quiet)
        pthread_join(display_thread, &res);

    if (transferts)
        free(transferts);

    printf("\n");

    return ret;
}