prolefeeder/prolefeeder.py

from flask import (Flask, render_template, flash, redirect,
                   url_for, request, send_file, abort)
from flask_bootstrap import Bootstrap
from flask_wtf import FlaskForm
from wtforms import DateTimeField, IntegerField, SubmitField
from wtforms.validators import InputRequired, NumberRange

from datetime import datetime, timedelta
from tempfile import mkstemp
import os
import random
import re
import subprocess

from config import Config


app = Flask(__name__)
app.config.from_object(Config)
Bootstrap(app)

attachment_filenames = {}

class DownloadForm(FlaskForm):
    start_time = DateTimeField('Start time',
                               validators=[InputRequired()],
                               format='%Y-%m-%d %H:%M')
    length = IntegerField('Length',
                          validators=[
                              InputRequired(),
                              NumberRange(min=1, max=app.config['MAX_LENGTH'])])
    submit = SubmitField('Download')


@app.route('/', methods=['GET', 'POST'])
def download():
    form = DownloadForm()
    if request.method == 'GET':
        form.start_time.data = datetime.now().replace(minute=0, second=0) - timedelta(hours=1)
        form.length.data = 60
    elif form.validate_on_submit():
        try:
            output_filename, attachment_filename = prepare_download(form)
            attachment_filenames[output_filename] = attachment_filename
            flash('The download should start immediately.', 'success')
            return render_template('download.html', form=form, filename=output_filename)
        except ValueError as e:
            flash('Error preparing download: {}'.format(e), 'warning')
    else:
        flash('Error in form!', 'warning')
    return render_template('download.html', form=form)


@app.route('/download_file/<filename>')
def download_file(filename):
    """Download an output file"""

    # Get attachment filename. This also makes sure that the user only downloads
    # (and removes) a file generated by us.
    try:
        attachment_filename = attachment_filenames.pop(filename)
    except KeyError:
        abort(404)

    fh = open(os.path.join(app.config['TMP_DIR'], filename), 'rb')
    os.remove(os.path.join(app.config['TMP_DIR'], filename))

    # XXX How to close fh?
    return send_file(fh, as_attachment=True,
                     attachment_filename=attachment_filename)


def prepare_download(form):
    """Prepare a download given the user's request form"""

    def start_time_for_source_fn(fn):
        return datetime.strptime(fn, 'qfhi-%Y%m%d-%H%M.mp3')

    def length_for_source_fn(fn):
        size = os.stat(os.path.join(app.config['DATA_DIR'], fn)).st_size
        length = timedelta(minutes=size / (1000*app.config['KBITS']/8) / 60)
        return length

    # Get a sorted list of all source files with start time and length
    sources = []
    for fn in os.listdir(app.config['DATA_DIR']):
        try:
            start_time = start_time_for_source_fn(fn)
            length = length_for_source_fn(fn)
            sources.append({'fn': fn, 'start_time': start_time, 'length': length})
        except ValueError:
            pass
    sources = sorted(sources, key=lambda s: s['start_time'])

    # Only interested in the source files from the start file
    start_index = None
    for i, source in enumerate(sources):
        if source['start_time'] <= form.start_time.data < source['start_time'] + source['length']:
            start_index = i
    if start_index is None:
        raise ValueError('Could not find start file')
    sources = sources[start_index:]

    # Super lazy: Limit to 5 source files input
    sources = sources[:5]

    # Seek into the first file
    ss = (form.start_time.data - sources[0]['start_time']).total_seconds()

    # Let ffmpeg do the rest of the job
    _, output_filename = mkstemp(suffix='.mp3', dir=app.config['TMP_DIR'])
    output_filename = os.path.basename(output_filename)
    attachment_filename = '{}_{}.mp3'.format(form.start_time.data, form.length.data)

    c = ['ffmpeg', '-y']
    c += ['-ss', str(ss)]
    c += ['-i', 'concat:' + '|'.join([os.path.join(app.config['DATA_DIR'],
        source['fn']) for source in sources])]
    c += ['-codec', 'copy']
    c += ['-t', str(form.length.data * 60)]
    c += [os.path.join(app.config['TMP_DIR'], output_filename)]

    app.logger.debug(' '.join(c))
    subprocess.call(c)

    return output_filename, attachment_filename


if __name__ == '__main__':
    app.run(debug=True)