<commit_before>#!/usr/bin/python
# -*- coding: utf8 -*-

"""
Audio library with conversion methods.

Author: Eduardo Ferreira
License: MIT (see LICENSE for details)
"""

# Module import
# ----------------------------------------------------------------------------------------------------------------------
from general import is_string_empty, update_extension, update_path

from enum import Enum
from json import dump, load
from os.path import basename, join
from re import match
from subprocess import call, CalledProcessError, check_output, PIPE, Popen
import logging
import sys

# Logger
# ----------------------------------------------------------------------------------------------------------------------
logging.basicConfig(level=logging.INFO)
_logger = logging.getLogger(__name__)

# Constants :: Error messages
# ----------------------------------------------------------------------------------------------------------------------
ERROR_PROGRAM_NOT_FOUND = 'Program \'{}\' was not found!'
"""
ERROR_WRONG_FILE_TYPE = 'The file \'{0}\' is not a valid {1} file!'
WARNING_NO_JSON_FILE = 'No JSON file with ID3 tags was found. Proceeding with encoding operation without tags.'
"""

# Constants :: Lists and file extensions
# ----------------------------------------------------------------------------------------------------------------------
ENCODING = 'utf-8'

"""
TITLE       : Track/Work name
ARTIST      : The artist generally considered responsible for the work. In popular music this is usually the performing
              band or singer. For classical music it would be the composer. For an audio book it would be the author of
              the original text
ALBUM       : The collection name to which this track belongs
TRACKNUMBER : The track number of this piece if part of a specific larger collection or album
ALBUMARTIST : The artist(s) who performed the work. In classical music this would be the conductor, orchestra, soloists.
              In an audio book it would be the actor who did the reading. In popular music this is typically the same as
              the ARTIST and is omitted
GENRE       : A short text indication of music genre
DATE        : Date the track was recorded
DISCTOTAL   : Number of discs that compose the set
TRACKTOTAL  : Number of tracks in the album
"""
TAGS = {
    'TITLE': '--tt',
    'ARTIST': '--ta',
    'ALBUM': '--tl',
    'TRACKNUMBER': '--tn',
    'ALBUMARTIST': ['--tv', 'TPE2='],
    'GENRE': '--tg',
    'DATE': '--ty',
    'DISCTOTAL': ['--tv', 'TPOS='],
    'TRACKTOTAL': ''
}

# Classes
# ----------------------------------------------------------------------------------------------------------------------
# Represents audio file extensions
class AudioFile(Enum):
    flac = '.flac'
    mp3 = '.mp3'
    wav = '.wav'

# Represents external programs used in the encoding and decoding operations
class Programs(Enum):
    flac = 'flac'
    metaflac = 'metaflac'
    lame = 'lame'
    grep = 'grep'
    sed = 'sed'
    file = 'file'

# Methods :: File encoding and decoding
# ----------------------------------------------------------------------------------------------------------------------
def decode_flac_wav(filename, destination, extract_cover=False, extract_tags=False):
    """
    Decodes a FLAC audio file, generating the corresponding WAV audio file.
    Also retrieves its ID3 tags and album cover.
    :param filename: The input audio file name
    :param destination: The destination where the output file will be stored
    :param extract_cover: Indicates if the album art should be extracted from the audio file
    :param extract_tags: Indicates if the ID3 tags should be extracted from the audio file
    :return: A tuple with three file names: output audio file, album art file and ID3 tags file
    """
    is_program_available(Programs.flac.value)

    if not is_flac_file(filename):
        return None

    # Invokes the 'flac' program with the following arguments:
    # -d => Decode (the default behavior is to encode)
    # -f => Force overwriting of output files
    # -o => Force the output file name
    output_filename = update_path(filename, destination, AudioFile.wav.value)
    call([Programs.flac.value, '-df', filename, '-o', output_filename])

    # Checks if both cover and tags should be retrieved
    cover = get_cover(filename, destination) if extract_cover else None
    tags = get_tags(filename) if extract_tags else None

    return output_filename, cover, tags

def encode_wav_flac(filename, destination, cover=None, tags=None):
    """
    Encodes a WAV audio file, generating the corresponding FLAC audio file.
    :param filename: The input audio file name
    :param destination: The destination where the output file will be stored
    :param cover: The name of the file with the album art
    :param tags: The name of the file with the ID3 tags
    :return: The name of the output audio file
    """
    is_program_available(Programs.flac.value)

    if not is_wav_file(filename):
        return None

    # Prepares the 'flac' program arguments:
    # -f => Force overwriting of output files
    # -8 => Synonymous with -l 12 -b 4096 -m -e -r 6
    # -V => Verify a correct encoding
    # -o => Force the output file name
    output_filename = update_path(filename, destination, AudioFile.flac.value)
    flac = [Programs.flac.value, '-f8V', '-o', output_filename]

    # Prepares the cover file to be passed as a parameter
    # --picture=SPECIFICATION => Import picture and store in PICTURE block
    if cover:
        flac.extend(['--picture=3||' + basename(cover) + '||' + cover])

    # Prepares the FLAC tags to be passed as parameters
    # --T FIELD=VALUE => Add a FLAC tag; may appear multiple times
    if tags:
        for tag, value in tags.items():
            flac.extend(['-T', tag + '=' + value])

    # Invokes the 'flac' program
    flac.append(filename)
    call(flac)

    return output_filename

def encode_wav_mp3(filename, destination, cover=None, tags=None):
    """
    Encodes a WAV audio file, generating the corresponding MP3 audio file.
    :param filename: The input audio file name
    :param destination: The destination where the output file will be stored
    :param cover: The name of the file with the album art
    :param tags: The name of the file with the ID3 tags
    :return: The name of the output audio file
    """
    is_program_available(Programs.lame.value)

    if not is_wav_file(filename):
        return None

    # Prepares the 'lame' program arguments:
    # -b 320          => Set the bitrate to 320 kbps
    # -q 0            => Highest quality, very slow
    # --preset insane => Type of the quality settings
    # --id3v2-only    => Add only a version 2 tag
    output_filename = update_path(filename, destination, AudioFile.mp3.value)
    lame = [Programs.lame.value, '-b', '320', '-q', '0', '--preset', 'insane', '--id3v2-only']

    # Prepares the cover file to be passed as a parameter
    # --ti <file> => Audio/song albumArt (jpeg/png/gif file, v2.3 tag)
    if cover:
        lame.extend(['--ti', cover])

    # Prepares the ID3 tags to be passed as parameters
    # --<tag> <value> => Audio/song specific information
    if tags:
        for tag, value in tags.items():
            id3_tag = TAGS[tag]
            if not id3_tag:
                continue

            if tag == 'TRACKNUMBER':
                value += '/' + tags['TRACKTOTAL']

            lame.extend([id3_tag, value] if type(id3_tag) is not list else [id3_tag[0], id3_tag[1] + value])

    # Invokes the 'lame' program
    lame.extend([filename, output_filename])
    call(lame)

    return output_filename

def encode_flac_mp3(filename, destination, extract_cover=False, extract_tags=False):
    """
    Decodes a FLAC audio file, generating the corresponding WAV audio file.
    The WAV audio file is then encoded, generating the corresponding MP3 audio file.
    :param filename: The input audio file name
    :param destination: The destination where the output file will be stored
    :param extract_cover: Indicates if the album art should be extracted from the audio file
    :param extract_tags: Indicates if the ID3 tags should be extracted from the audio file
    :return: The name of the output audio file
    """
    wav_file = decode_flac_wav(filename, destination, extract_cover, extract_tags)
    if wav_file:
        cover = wav_file[1] if get_cover else None
        tags = wav_file[2] if get_tags else None
        return encode_wav_mp3(wav_file[0], destination, cover, tags)

    return None

# Methods :: Album cover management
# ----------------------------------------------------------------------------------------------------------------------
def get_cover(filename, destination):
    """
    Retrieves the front cover art file from a FLAC audio file and stores it in the destination directory.
    :param filename: The input audio file name
    :param destination: The destination where the output file will be stored
    ;return: The name of the album art file
    """
    is_program_available(Programs.metaflac.value)
    is_program_available(Programs.grep.value)
    is_program_available(Programs.sed.value)

    # Prepares the 'metaflac' program arguments:
    # --list       => Lists the full stack of metadata
    # --block-type => Comma-separated list of block types to be included
    metaflac = [Programs.metaflac.value, '--list', '--block-type=PICTURE', filename]

    # Prepares the 'grep' program arguments (looks for description parameter)
    grep = [Programs.grep.value, 'description:']

    # Prepares the 'sed' program arguments (the regular expression removes the parameter name)
    sed = [Programs.sed.value, 's/.*: //']

    # Invokes the 'metaflac', 'grep' and 'sed' programs to retrieve the cover file name
    p1 = Popen(metaflac, stdout=PIPE)
    p2 = Popen(grep, stdin=p1.stdout, stdout=PIPE)
    p3 = Popen(sed, stdin=p2.stdout, stdout=PIPE)
    cover = p3.stdout.read().decode(ENCODING).rstrip('\n')

    # Checks if the audio file has a cover
    if is_string_empty(cover):
        return None

    # Prepares the 'metaflac' program arguments:
    # --export-picture-to => Export PICTURE block to a file
    cover = join(destination, cover)
    call([Programs.metaflac.value, '--export-picture-to=' + cover, filename])

    return cover

# Methods :: Tag management
# ----------------------------------------------------------------------------------------------------------------------
def read_tags(filename):
    """
    Reads a JSON file with ID3 tags.
    :param filename: The input audio file name
    :return: The name of the ID3 tags file
    """
    try:
        with open(update_extension(filename, '.json'), 'r') as tags_file:
            tags = load(tags_file)
    except FileNotFoundError:
        return None

    return tags

def write_tags(filename, tags):
    """
    Writes ID3 tags to a JSON file.
    :param filename: The input audio file name
    :param tags: The list of ID3 tags written to the audio file
    """
    with open(update_extension(filename, '.json'), 'w') as tags_file:
        dump(tags, tags_file, indent=4)

def get_tags(filename):
    """
    Retrieves the tag values of a FLAC audio file.
    :param filename: The input audio file name
    :return: The list of ID3 tags retrieved from the audio file
    """
    is_program_available(Programs.metaflac.value)

    tags = {}
    for tag in TAGS:

        # Invokes the 'metaflac' program with the following arguments:
        # --show-tag => Shows the value of the given tag
        p = Popen([Programs.metaflac.value, '--show-tag=' + tag, filename], stdout=PIPE)
        value = p.stdout.read().decode(ENCODING).rstrip('\n')
        if value:
            tags[tag] = value.split('=')[1]

    return tags

# Methods :: File management
# ----------------------------------------------------------------------------------------------------------------------
def is_flac_file(filename):
    """
    Checks if the given file is a valid FLAC audio file.
    :param filename: The input audio file name
    :return: True if the input file is a FLAC audio file; False otherwise
    """
    is_program_available(Programs.metaflac.value)

    # Invokes the 'metaflac' program with the following arguments:
    # --show-md5sum => Show the MD5 signature from the STREAMINFO block
    output = ''
    try:
        output = check_output([Programs.metaflac.value, '--show-md5sum', filename]).decode(ENCODING)
    except CalledProcessError as e:
        if e.returncode == 1:
            return False

    return match(r'[a-z0-9]+', output)

def is_wav_file(filename):
    """
    Checks if the given file is a valid WAV audio file.
    :param filename: The input audio file name
    :return: True if the input file is a WAV audio file; False otherwise
    """
    is_program_available(Programs.file.value)

    # Invokes the 'file' program with the following arguments:
    # --mime-type => Output the MIME type
    output = ''
    try:
        output = check_output([Programs.file.value, '--mime-type', filename]).decode(ENCODING)
    except CalledProcessError as e:
        if e.returncode == 1:
            return False

    return 'audio/x-wav' in output

# Methods :: External program validation
# ----------------------------------------------------------------------------------------------------------------------
def is_program_available(program):
    """
    Checks if an external program is present in the operating system.
    :param program: The name of the external program
    :return: True if the program is present in the operating system; False otherwise
    """
    # The output of the following command should be something like:
    # flac: /usr/bin/flac /usr/share/man/man1/flac.1.gz
    output = Popen(['whereis', program], stdout=PIPE).communicate()[0]
    if len(output.split()) == 1:
        _logger.error(ERROR_PROGRAM_NOT_FOUND.format(program))
        sys.exit(1)
<commit_msg>Reduced the number of columns in the source code to 100.<commit_after>#!/usr/bin/python
# -*- coding: utf8 -*-

"""
Audio library with conversion methods.

Author: Eduardo Ferreira
License: MIT (see LICENSE for details)
"""

# Module import
# --------------------------------------------------------------------------------------------------
from general import is_string_empty, update_extension, update_path

from enum import Enum
from json import dump, load
from os.path import basename, join
from re import match
from subprocess import call, CalledProcessError, check_output, PIPE, Popen
import logging
import sys

# Logger
# --------------------------------------------------------------------------------------------------
logging.basicConfig(level=logging.INFO)
_logger = logging.getLogger(__name__)

# Constants :: Error messages
# --------------------------------------------------------------------------------------------------
ERROR_PROGRAM_NOT_FOUND = 'Program \'{}\' was not found!'

# Constants :: Lists and file extensions
# --------------------------------------------------------------------------------------------------
ENCODING = 'utf-8'

"""
TITLE       : Track/Work name
ARTIST      : The artist generally considered responsible for the work. In popular music this is
              usually the performing band or singer. For classical music it would be the composer.
              For an audio book it would be the author of the original text
ALBUM       : The collection name to which this track belongs
TRACKNUMBER : The track number of this piece if part of a specific larger collection or album
ALBUMARTIST : The artist(s) who performed the work. In classical music this would be the conductor,
              orchestra, soloists. In an audio book it would be the actor who did the reading. In
              popular music this is typically the same as the ARTIST and is omitted
GENRE       : A short text indication of music genre
DATE        : Date the track was recorded
DISCTOTAL   : Number of discs that compose the set
TRACKTOTAL  : Number of tracks in the album
"""
TAGS = {
    'TITLE': '--tt',
    'ARTIST': '--ta',
    'ALBUM': '--tl',
    'TRACKNUMBER': '--tn',
    'ALBUMARTIST': ['--tv', 'TPE2='],
    'GENRE': '--tg',
    'DATE': '--ty',
    'DISCTOTAL': ['--tv', 'TPOS='],
    'TRACKTOTAL': ''
}

# Classes
# --------------------------------------------------------------------------------------------------
# Represents audio file extensions
class AudioFile(Enum):
    flac = '.flac'
    mp3 = '.mp3'
    wav = '.wav'

# Represents external programs used in the encoding and decoding operations
class Programs(Enum):
    flac = 'flac'
    metaflac = 'metaflac'
    lame = 'lame'
    grep = 'grep'
    sed = 'sed'
    file = 'file'

# Methods :: File encoding and decoding
# --------------------------------------------------------------------------------------------------
def decode_flac_wav(filename, destination, extract_cover=False, extract_tags=False):
    """
    Decodes a FLAC audio file, generating the corresponding WAV audio file.
    Also retrieves its ID3 tags and album cover.
    :param filename: The input audio file name
    :param destination: The destination where the output file will be stored
    :param extract_cover: Indicates if the album art should be extracted from the audio file
    :param extract_tags: Indicates if the ID3 tags should be extracted from the audio file
    :return: A tuple with three file names: output audio file, album art file and ID3 tags file
    """
    is_program_available(Programs.flac.value)

    if not is_flac_file(filename):
        return None

    # Invokes the 'flac' program with the following arguments:
    # -d => Decode (the default behavior is to encode)
    # -f => Force overwriting of output files
    # -o => Force the output file name
    output_filename = update_path(filename, destination, AudioFile.wav.value)
    call([Programs.flac.value, '-df', filename, '-o', output_filename])

    # Checks if both cover and tags should be retrieved
    cover = get_cover(filename, destination) if extract_cover else None
    tags = get_tags(filename) if extract_tags else None

    return output_filename, cover, tags

def encode_wav_flac(filename, destination, cover=None, tags=None):
    """
    Encodes a WAV audio file, generating the corresponding FLAC audio file.
    :param filename: The input audio file name
    :param destination: The destination where the output file will be stored
    :param cover: The name of the file with the album art
    :param tags: The name of the file with the ID3 tags
    :return: The name of the output audio file
    """
    is_program_available(Programs.flac.value)

    if not is_wav_file(filename):
        return None

    # Prepares the 'flac' program arguments:
    # -f => Force overwriting of output files
    # -8 => Synonymous with -l 12 -b 4096 -m -e -r 6
    # -V => Verify a correct encoding
    # -o => Force the output file name
    output_filename = update_path(filename, destination, AudioFile.flac.value)
    flac = [Programs.flac.value, '-f8V', '-o', output_filename]

    # Prepares the cover file to be passed as a parameter
    # --picture=SPECIFICATION => Import picture and store in PICTURE block
    if cover:
        flac.extend(['--picture=3||' + basename(cover) + '||' + cover])

    # Prepares the FLAC tags to be passed as parameters
    # --T FIELD=VALUE => Add a FLAC tag; may appear multiple times
    if tags:
        for tag, value in tags.items():
            flac.extend(['-T', tag + '=' + value])

    # Invokes the 'flac' program
    flac.append(filename)
    call(flac)

    return output_filename

def encode_wav_mp3(filename, destination, cover=None, tags=None):
    """
    Encodes a WAV audio file, generating the corresponding MP3 audio file.
    :param filename: The input audio file name
    :param destination: The destination where the output file will be stored
    :param cover: The name of the file with the album art
    :param tags: The name of the file with the ID3 tags
    :return: The name of the output audio file
    """
    is_program_available(Programs.lame.value)

    if not is_wav_file(filename):
        return None

    # Prepares the 'lame' program arguments:
    # -b 320          => Set the bitrate to 320 kbps
    # -q 0            => Highest quality, very slow
    # --preset insane => Type of the quality settings
    # --id3v2-only    => Add only a version 2 tag
    output_filename = update_path(filename, destination, AudioFile.mp3.value)
    lame = [Programs.lame.value, '-b', '320', '-q', '0', '--preset', 'insane', '--id3v2-only']

    # Prepares the cover file to be passed as a parameter
    # --ti <file> => Audio/song albumArt (jpeg/png/gif file, v2.3 tag)
    if cover:
        lame.extend(['--ti', cover])

    # Prepares the ID3 tags to be passed as parameters
    # --<tag> <value> => Audio/song specific information
    if tags:
        for tag, value in tags.items():
            id3_tag = TAGS[tag]
            if not id3_tag:
                continue

            if tag == 'TRACKNUMBER':
                value += '/' + tags['TRACKTOTAL']

            lame.extend([id3_tag, value]
                if type(id3_tag) is not list
                else [id3_tag[0], id3_tag[1] + value])

    # Invokes the 'lame' program
    lame.extend([filename, output_filename])
    call(lame)

    return output_filename

def encode_flac_mp3(filename, destination, extract_cover=False, extract_tags=False):
    """
    Decodes a FLAC audio file, generating the corresponding WAV audio file.
    The WAV audio file is then encoded, generating the corresponding MP3 audio file.
    :param filename: The input audio file name
    :param destination: The destination where the output file will be stored
    :param extract_cover: Indicates if the album art should be extracted from the audio file
    :param extract_tags: Indicates if the ID3 tags should be extracted from the audio file
    :return: The name of the output audio file
    """
    wav_file = decode_flac_wav(filename, destination, extract_cover, extract_tags)
    if wav_file:
        cover = wav_file[1] if get_cover else None
        tags = wav_file[2] if get_tags else None
        return encode_wav_mp3(wav_file[0], destination, cover, tags)

    return None

# Methods :: Album cover management
# --------------------------------------------------------------------------------------------------
def get_cover(filename, destination):
    """
    Retrieves the front cover art file from a FLAC audio file and stores it in the destination
    directory.
    :param filename: The input audio file name
    :param destination: The destination where the output file will be stored
    ;return: The name of the album art file
    """
    is_program_available(Programs.metaflac.value)
    is_program_available(Programs.grep.value)
    is_program_available(Programs.sed.value)

    # Prepares the 'metaflac' program arguments:
    # --list       => Lists the full stack of metadata
    # --block-type => Comma-separated list of block types to be included
    metaflac = [Programs.metaflac.value, '--list', '--block-type=PICTURE', filename]

    # Prepares the 'grep' program arguments (looks for description parameter)
    grep = [Programs.grep.value, 'description:']

    # Prepares the 'sed' program arguments (the regular expression removes the parameter name)
    sed = [Programs.sed.value, 's/.*: //']

    # Invokes the 'metaflac', 'grep' and 'sed' programs to retrieve the cover file name
    p1 = Popen(metaflac, stdout=PIPE)
    p2 = Popen(grep, stdin=p1.stdout, stdout=PIPE)
    p3 = Popen(sed, stdin=p2.stdout, stdout=PIPE)
    cover = p3.stdout.read().decode(ENCODING).rstrip('\n')

    # Checks if the audio file has a cover
    if is_string_empty(cover):
        return None

    # Prepares the 'metaflac' program arguments:
    # --export-picture-to => Export PICTURE block to a file
    cover = join(destination, cover)
    call([Programs.metaflac.value, '--export-picture-to=' + cover, filename])

    return cover

# Methods :: Tag management
# --------------------------------------------------------------------------------------------------
def read_tags(filename):
    """
    Reads a JSON file with ID3 tags.
    :param filename: The input audio file name
    :return: The name of the ID3 tags file
    """
    try:
        with open(update_extension(filename, '.json'), 'r') as tags_file:
            tags = load(tags_file)
    except FileNotFoundError:
        return None

    return tags

def write_tags(filename, tags):
    """
    Writes ID3 tags to a JSON file.
    :param filename: The input audio file name
    :param tags: The list of ID3 tags written to the audio file
    """
    with open(update_extension(filename, '.json'), 'w') as tags_file:
        dump(tags, tags_file, indent=4)

def get_tags(filename):
    """
    Retrieves the tag values of a FLAC audio file.
    :param filename: The input audio file name
    :return: The list of ID3 tags retrieved from the audio file
    """
    is_program_available(Programs.metaflac.value)

    tags = {}
    for tag in TAGS:

        # Invokes the 'metaflac' program with the following arguments:
        # --show-tag => Shows the value of the given tag
        p = Popen([Programs.metaflac.value, '--show-tag=' + tag, filename], stdout=PIPE)
        value = p.stdout.read().decode(ENCODING).rstrip('\n')
        if value:
            tags[tag] = value.split('=')[1]

    return tags

# Methods :: File management
# --------------------------------------------------------------------------------------------------
def is_flac_file(filename):
    """
    Checks if the given file is a valid FLAC audio file.
    :param filename: The input audio file name
    :return: True if the input file is a FLAC audio file; False otherwise
    """
    is_program_available(Programs.metaflac.value)

    # Invokes the 'metaflac' program with the following arguments:
    # --show-md5sum => Show the MD5 signature from the STREAMINFO block
    output = ''
    try:
        output = check_output([Programs.metaflac.value, '--show-md5sum', filename]).decode(ENCODING)
    except CalledProcessError as e:
        if e.returncode == 1:
            return False

    return match(r'[a-z0-9]+', output)

def is_wav_file(filename):
    """
    Checks if the given file is a valid WAV audio file.
    :param filename: The input audio file name
    :return: True if the input file is a WAV audio file; False otherwise
    """
    is_program_available(Programs.file.value)

    # Invokes the 'file' program with the following arguments:
    # --mime-type => Output the MIME type
    output = ''
    try:
        output = check_output([Programs.file.value, '--mime-type', filename]).decode(ENCODING)
    except CalledProcessError as e:
        if e.returncode == 1:
            return False

    return 'audio/x-wav' in output

# Methods :: External program validation
# --------------------------------------------------------------------------------------------------
def is_program_available(program):
    """
    Checks if an external program is present in the operating system.
    :param program: The name of the external program
    :return: True if the program is present in the operating system; False otherwise
    """
    # The output of the following command should be something like:
    # flac: /usr/bin/flac /usr/share/man/man1/flac.1.gz
    output = Popen(['whereis', program], stdout=PIPE).communicate()[0]
    if len(output.split()) == 1:
        _logger.error(ERROR_PROGRAM_NOT_FOUND.format(program))
        sys.exit(1)
<|endoftext|>


<commit_before>from sqlalchemy import Column, Integer, String, Boolean
from database import Base

# IMPORTANT!
# VALID MODULES:

# applications
# library         [mandatory static]
# recently_added
# sabnzbd
# synopsis        [mandatory static]
# trakt           [mandatory static]

class Module(Base):
    __tablename__ = 'modules'
    id = Column(Integer, primary_key=True)
    name = Column(String(50))
    static = Column(Boolean)
    column = Column(Integer)
    position = Column(Integer)
    poll = Column(Integer)
    delay = Column(Integer)

    def __init__(self, name=None, static=0, column=None, position=None, poll=None, delay=None):
        self.name = name
        self.static = static
        self.column = column
        self.position = position
        self.poll = poll
        self.delay = delay

    def __repr__(self):
        return '<Module %r>' % (self.name)
<commit_msg>should only be able to have 1 of each module<commit_after>from sqlalchemy import Column, Integer, String, Boolean
from database import Base

# IMPORTANT!
# VALID MODULES:

# applications
# library         [mandatory static]
# recently_added
# sabnzbd
# synopsis        [mandatory static]
# trakt           [mandatory static]

class Module(Base):
    __tablename__ = 'modules'
    id = Column(Integer, primary_key=True)
    name = Column(String(50), unique=True)
    static = Column(Boolean)
    column = Column(Integer)
    position = Column(Integer)
    poll = Column(Integer)
    delay = Column(Integer)

    def __init__(self, name=None, static=0, column=None, position=None, poll=None, delay=None):
        self.name = name
        self.static = static
        self.column = column
        self.position = position
        self.poll = poll
        self.delay = delay

    def __repr__(self):
        return '<Module %r>' % (self.name)
<|endoftext|>


<commit_before>"""
A collection of models we'll use to attempt to classify videos.
"""
from keras.layers import Dense, Flatten, Dropout, ZeroPadding3D
from keras.layers.recurrent import LSTM
from keras.models import Sequential, load_model
from keras.optimizers import Adam
from keras.layers.wrappers import TimeDistributed
from keras.layers.convolutional import (Conv2D, MaxPooling3D, Conv3D,
    MaxPooling2D)
from collections import deque
import sys

class ResearchModels():
    def __init__(self, nb_classes, model, seq_length,
                 saved_model=None, features_length=2048):
        """
        `model` = one of:
            lstm
            crnn
            mlp
            conv_3d
        `nb_classes` = the number of classes to predict
        `seq_length` = the length of our video sequences
        `saved_model` = the path to a saved Keras model to load
        """

        # Set defaults.
        self.seq_length = seq_length
        self.load_model = load_model
        self.saved_model = saved_model
        self.nb_classes = nb_classes
        self.feature_queue = deque()

        # Set the metrics. Only use top k if there's a need.
        metrics = ['accuracy']
        if self.nb_classes >= 10:
            metrics.append('top_k_categorical_accuracy')

        # Get the appropriate model.
        if self.saved_model is not None:
            print("Loading model %s" % self.saved_model)
            self.model = load_model(self.saved_model)
        elif model == 'lstm':
            print("Loading LSTM model.")
            self.input_shape = (seq_length, features_length)
            self.model = self.lstm()
        elif model == 'crnn':
            print("Loading CRNN model.")
            self.input_shape = (seq_length, 80, 80, 3)
            self.model = self.crnn()
        elif model == 'mlp':
            print("Loading simple MLP.")
            self.input_shape = features_length * seq_length
            self.model = self.mlp()
        elif model == 'conv_3d':
            print("Loading Conv3D")
            self.input_shape = (seq_length, 80, 80, 3)
            self.model = self.conv_3d()
        elif model == 'c3d':
            print("Loading C3D")
            self.input_shape = (seq_length, 80, 80, 3)
            self.model = self.c3d()
        else:
            print("Unknown network.")
            sys.exit()

        # Now compile the network.
        optimizer = Adam(lr=1e-5)
        self.model.compile(loss='categorical_crossentropy', optimizer=optimizer,
                           metrics=metrics)

    def lstm(self):
        """Build a simple LSTM network. We pass the extracted features from
        our CNN to this model predomenently."""
        # Model.
        model = Sequential()
        model.add(LSTM(2048, return_sequences=True, input_shape=self.input_shape,
                       dropout=0.5))
        model.add(Flatten())
        model.add(Dense(512, activation='relu'))
        model.add(Dropout(0.5))
        model.add(Dense(self.nb_classes, activation='softmax'))

        return model

    def crnn(self):
        """Build a CNN into RNN.
        Starting version from:
        https://github.com/udacity/self-driving-car/blob/master/
            steering-models/community-models/chauffeur/models.py
        """
        model = Sequential()
        model.add(TimeDistributed(Conv2D(32, (3,3),
            kernel_initializer="he_normal",
            activation='relu'), input_shape=self.input_shape))
        model.add(TimeDistributed(Conv2D(32, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(MaxPooling2D()))
        model.add(TimeDistributed(Conv2D(48, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(Conv2D(48, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(MaxPooling2D()))
        model.add(TimeDistributed(Conv2D(64, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(Conv2D(64, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(MaxPooling2D()))
        model.add(TimeDistributed(Conv2D(128, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(Conv2D(128, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(MaxPooling2D()))
        model.add(TimeDistributed(Flatten()))
        model.add(LSTM(256, return_sequences=True))
        model.add(Flatten())
        model.add(Dense(512))
        model.add(Dropout(0.5))
        model.add(Dense(self.nb_classes, activation='softmax'))

        return model

    def mlp(self):
        """Build a simple MLP."""
        # Model.
        model = Sequential()
        model.add(Dense(512, input_dim=self.input_shape))
        model.add(Dropout(0.5))
        model.add(Dense(512))
        model.add(Dropout(0.5))
        model.add(Dense(self.nb_classes, activation='softmax'))

        return model

    def conv_3d(self):
        """
        Build a 3D convolutional network, based loosely on C3D.
            https://arxiv.org/pdf/1412.0767.pdf
        """
        # Model.
        model = Sequential()
        model.add(Conv3D(
            32, (7,7,7), activation='relu', input_shape=self.input_shape
        ))
        model.add(MaxPooling3D(pool_size=(1, 2, 2), strides=(1, 2, 2)))
        model.add(Conv3D(64, (3,3,3), activation='relu'))
        model.add(MaxPooling3D(pool_size=(1, 2, 2), strides=(1, 2, 2)))
        model.add(Conv3D(128, (2,2,2), activation='relu'))
        model.add(MaxPooling3D(pool_size=(1, 2, 2), strides=(1, 2, 2)))
        model.add(Flatten())
        model.add(Dense(256))
        model.add(Dropout(0.2))
        model.add(Dense(256))
        model.add(Dropout(0.2))
        model.add(Dense(self.nb_classes, activation='softmax'))

        return model

    def c3d(self):
        """
        Build a 3D convolutional network, aka C3D.
            https://arxiv.org/pdf/1412.0767.pdf

        With thanks:
            https://gist.github.com/albertomontesg/d8b21a179c1e6cca0480ebdf292c34d2
        """
        model = Sequential()
        # 1st layer group
        model.add(Conv3D(64, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv1',
                         subsample=(1, 1, 1),
                         input_shape=self.input_shape))
        model.add(MaxPooling3D(pool_size=(1, 2, 2), strides=(1, 2, 2),
                               border_mode='valid', name='pool1'))
        # 2nd layer group
        model.add(Conv3D(128, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv2',
                         subsample=(1, 1, 1)))
        model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2),
                               border_mode='valid', name='pool2'))
        # 3rd layer group
        model.add(Conv3D(256, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv3a',
                         subsample=(1, 1, 1)))
        model.add(Conv3D(256, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv3b',
                         subsample=(1, 1, 1)))
        model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2),
                               border_mode='valid', name='pool3'))
        # 4th layer group
        model.add(Conv3D(512, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv4a',
                         subsample=(1, 1, 1)))
        model.add(Conv3D(512, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv4b',
                         subsample=(1, 1, 1)))
        model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2),
                               border_mode='valid', name='pool4'))

        # 5th layer group
        model.add(Conv3D(512, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv5a',
                         subsample=(1, 1, 1)))
        model.add(Conv3D(512, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv5b',
                         subsample=(1, 1, 1)))
        model.add(ZeroPadding3D(padding=(0, 1, 1)))
        model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2),
                               border_mode='valid', name='pool5'))
        model.add(Flatten())

        # FC layers group
        model.add(Dense(4096, activation='relu', name='fc6'))
        model.add(Dropout(0.5))
        model.add(Dense(4096, activation='relu', name='fc7'))
        model.add(Dropout(0.5))
        model.add(Dense(self.nb_classes, activation='softmax'))
<commit_msg>Return the model<commit_after>"""
A collection of models we'll use to attempt to classify videos.
"""
from keras.layers import Dense, Flatten, Dropout, ZeroPadding3D
from keras.layers.recurrent import LSTM
from keras.models import Sequential, load_model
from keras.optimizers import Adam
from keras.layers.wrappers import TimeDistributed
from keras.layers.convolutional import (Conv2D, MaxPooling3D, Conv3D,
    MaxPooling2D)
from collections import deque
import sys

class ResearchModels():
    def __init__(self, nb_classes, model, seq_length,
                 saved_model=None, features_length=2048):
        """
        `model` = one of:
            lstm
            crnn
            mlp
            conv_3d
        `nb_classes` = the number of classes to predict
        `seq_length` = the length of our video sequences
        `saved_model` = the path to a saved Keras model to load
        """

        # Set defaults.
        self.seq_length = seq_length
        self.load_model = load_model
        self.saved_model = saved_model
        self.nb_classes = nb_classes
        self.feature_queue = deque()

        # Set the metrics. Only use top k if there's a need.
        metrics = ['accuracy']
        if self.nb_classes >= 10:
            metrics.append('top_k_categorical_accuracy')

        # Get the appropriate model.
        if self.saved_model is not None:
            print("Loading model %s" % self.saved_model)
            self.model = load_model(self.saved_model)
        elif model == 'lstm':
            print("Loading LSTM model.")
            self.input_shape = (seq_length, features_length)
            self.model = self.lstm()
        elif model == 'crnn':
            print("Loading CRNN model.")
            self.input_shape = (seq_length, 80, 80, 3)
            self.model = self.crnn()
        elif model == 'mlp':
            print("Loading simple MLP.")
            self.input_shape = features_length * seq_length
            self.model = self.mlp()
        elif model == 'conv_3d':
            print("Loading Conv3D")
            self.input_shape = (seq_length, 80, 80, 3)
            self.model = self.conv_3d()
        elif model == 'c3d':
            print("Loading C3D")
            self.input_shape = (seq_length, 80, 80, 3)
            self.model = self.c3d()
        else:
            print("Unknown network.")
            sys.exit()

        # Now compile the network.
        optimizer = Adam(lr=1e-5)
        self.model.compile(loss='categorical_crossentropy', optimizer=optimizer,
                           metrics=metrics)

    def lstm(self):
        """Build a simple LSTM network. We pass the extracted features from
        our CNN to this model predomenently."""
        # Model.
        model = Sequential()
        model.add(LSTM(2048, return_sequences=True, input_shape=self.input_shape,
                       dropout=0.5))
        model.add(Flatten())
        model.add(Dense(512, activation='relu'))
        model.add(Dropout(0.5))
        model.add(Dense(self.nb_classes, activation='softmax'))

        return model

    def crnn(self):
        """Build a CNN into RNN.
        Starting version from:
        https://github.com/udacity/self-driving-car/blob/master/
            steering-models/community-models/chauffeur/models.py
        """
        model = Sequential()
        model.add(TimeDistributed(Conv2D(32, (3,3),
            kernel_initializer="he_normal",
            activation='relu'), input_shape=self.input_shape))
        model.add(TimeDistributed(Conv2D(32, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(MaxPooling2D()))
        model.add(TimeDistributed(Conv2D(48, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(Conv2D(48, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(MaxPooling2D()))
        model.add(TimeDistributed(Conv2D(64, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(Conv2D(64, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(MaxPooling2D()))
        model.add(TimeDistributed(Conv2D(128, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(Conv2D(128, (3,3),
            kernel_initializer="he_normal",
            activation='relu')))
        model.add(TimeDistributed(MaxPooling2D()))
        model.add(TimeDistributed(Flatten()))
        model.add(LSTM(256, return_sequences=True))
        model.add(Flatten())
        model.add(Dense(512))
        model.add(Dropout(0.5))
        model.add(Dense(self.nb_classes, activation='softmax'))

        return model

    def mlp(self):
        """Build a simple MLP."""
        # Model.
        model = Sequential()
        model.add(Dense(512, input_dim=self.input_shape))
        model.add(Dropout(0.5))
        model.add(Dense(512))
        model.add(Dropout(0.5))
        model.add(Dense(self.nb_classes, activation='softmax'))

        return model

    def conv_3d(self):
        """
        Build a 3D convolutional network, based loosely on C3D.
            https://arxiv.org/pdf/1412.0767.pdf
        """
        # Model.
        model = Sequential()
        model.add(Conv3D(
            32, (7,7,7), activation='relu', input_shape=self.input_shape
        ))
        model.add(MaxPooling3D(pool_size=(1, 2, 2), strides=(1, 2, 2)))
        model.add(Conv3D(64, (3,3,3), activation='relu'))
        model.add(MaxPooling3D(pool_size=(1, 2, 2), strides=(1, 2, 2)))
        model.add(Conv3D(128, (2,2,2), activation='relu'))
        model.add(MaxPooling3D(pool_size=(1, 2, 2), strides=(1, 2, 2)))
        model.add(Flatten())
        model.add(Dense(256))
        model.add(Dropout(0.2))
        model.add(Dense(256))
        model.add(Dropout(0.2))
        model.add(Dense(self.nb_classes, activation='softmax'))

        return model

    def c3d(self):
        """
        Build a 3D convolutional network, aka C3D.
            https://arxiv.org/pdf/1412.0767.pdf

        With thanks:
            https://gist.github.com/albertomontesg/d8b21a179c1e6cca0480ebdf292c34d2
        """
        model = Sequential()
        # 1st layer group
        model.add(Conv3D(64, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv1',
                         subsample=(1, 1, 1),
                         input_shape=self.input_shape))
        model.add(MaxPooling3D(pool_size=(1, 2, 2), strides=(1, 2, 2),
                               border_mode='valid', name='pool1'))
        # 2nd layer group
        model.add(Conv3D(128, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv2',
                         subsample=(1, 1, 1)))
        model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2),
                               border_mode='valid', name='pool2'))
        # 3rd layer group
        model.add(Conv3D(256, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv3a',
                         subsample=(1, 1, 1)))
        model.add(Conv3D(256, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv3b',
                         subsample=(1, 1, 1)))
        model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2),
                               border_mode='valid', name='pool3'))
        # 4th layer group
        model.add(Conv3D(512, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv4a',
                         subsample=(1, 1, 1)))
        model.add(Conv3D(512, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv4b',
                         subsample=(1, 1, 1)))
        model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2),
                               border_mode='valid', name='pool4'))

        # 5th layer group
        model.add(Conv3D(512, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv5a',
                         subsample=(1, 1, 1)))
        model.add(Conv3D(512, 3, 3, 3, activation='relu',
                         border_mode='same', name='conv5b',
                         subsample=(1, 1, 1)))
        model.add(ZeroPadding3D(padding=(0, 1, 1)))
        model.add(MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2),
                               border_mode='valid', name='pool5'))
        model.add(Flatten())

        # FC layers group
        model.add(Dense(4096, activation='relu', name='fc6'))
        model.add(Dropout(0.5))
        model.add(Dense(4096, activation='relu', name='fc7'))
        model.add(Dropout(0.5))
        model.add(Dense(self.nb_classes, activation='softmax'))

        return model
<|endoftext|>


<commit_before>import datetime

import mongoengine
from mongoengine.django import auth

from piplmesh.account import fields

LOWER_DATE_LIMIT = 366 * 120

POST_MESSAGE_MAX_LENGTH = 500
COMMENT_MESSAGE_MAX_LENGTH = 300

class User(auth.User):
    birthdate = fields.LimitedDateTimeField(upper_limit=datetime.datetime.today(), lower_limit=datetime.datetime.today() - datetime.timedelta(LOWER_DATE_LIMIT))
    gender = fields.GenderField()
    language = fields.LanguageField()

    facebook_id = mongoengine.IntField()
    facebook_token = mongoengine.StringField(max_length=150)

class Comment(mongoengine.EmbeddedDocument):
    """
    This class defines document type for comments on wall posts.
    """

    created_time = mongoengine.DateTimeField(default=lambda:datetime.datetime.now())
    author = mongoengine.ReferenceField(User, required=True)
    message = mongoengine.StringField(max_length=COMMENT_MESSAGE_MAX_LENGTH)

class Post(mongoengine.Document):
    """
    This class defines document type for posts on a wall.
    """

    author = mongoengine.ReferenceField(User, required=True)
    created_time = mongoengine.DateTimeField(default=lambda:datetime.datetime.now())
    updated_time = mongoengine.DateTimeField()
    comments = mongoengine.ListField(mongoengine.EmbeddedDocumentField(Comment))

class TextPost(Post):
    """
    This class defines document type for text posts.
    """

    message = mongoengine.StringField(max_length=POST_MESSAGE_MAX_LENGTH)

class ImagePost(Post):
    """
    This class defines document type for image posts.
    """

    image_path = mongoengine.StringField()

class LinkPost(Post):
    """
    This class defines document type for link posts.
    """

    link_url = mongoengine.URLField()
    link_caption = mongoengine.StringField()
    link_description = mongoengine.StringField()<commit_msg>Fixed space after comment.<commit_after>import datetime

import mongoengine
from mongoengine.django import auth

from piplmesh.account import fields

LOWER_DATE_LIMIT = 366 * 120

POST_MESSAGE_MAX_LENGTH = 500
COMMENT_MESSAGE_MAX_LENGTH = 300

class User(auth.User):
    birthdate = fields.LimitedDateTimeField(upper_limit=datetime.datetime.today(), lower_limit=datetime.datetime.today() - datetime.timedelta(LOWER_DATE_LIMIT))
    gender = fields.GenderField()
    language = fields.LanguageField()

    facebook_id = mongoengine.IntField()
    facebook_token = mongoengine.StringField(max_length=150)

class Comment(mongoengine.EmbeddedDocument):
    """
    This class defines document type for comments on wall posts.
    """

    created_time = mongoengine.DateTimeField(default=lambda: datetime.datetime.now())
    author = mongoengine.ReferenceField(User, required=True)
    message = mongoengine.StringField(max_length=COMMENT_MESSAGE_MAX_LENGTH)

class Post(mongoengine.Document):
    """
    This class defines document type for posts on a wall.
    """

    author = mongoengine.ReferenceField(User, required=True)
    created_time = mongoengine.DateTimeField(default=lambda: datetime.datetime.now())
    updated_time = mongoengine.DateTimeField()
    comments = mongoengine.ListField(mongoengine.EmbeddedDocumentField(Comment))

class TextPost(Post):
    """
    This class defines document type for text posts.
    """

    message = mongoengine.StringField(max_length=POST_MESSAGE_MAX_LENGTH)

class ImagePost(Post):
    """
    This class defines document type for image posts.
    """

    image_path = mongoengine.StringField()

class LinkPost(Post):
    """
    This class defines document type for link posts.
    """

    link_url = mongoengine.URLField()
    link_caption = mongoengine.StringField()
    link_description = mongoengine.StringField()<|endoftext|>


<commit_before>"""
websocket - WebSocket client library for Python

Copyright (C) 2010 Hiroki Ohtani(liris)

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

    This library 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
    Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General Public
    License along with this library; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

"""

import socket
from urlparse import urlparse
import random
import struct
import hashlib
import logging

logger = logging.getLogger()

class WebSocketException(Exception):
    pass

class ConnectionClosedException(WebSocketException):
    pass

default_timeout = None
traceEnabled = False

def enableTrace(tracable):
    """
    turn on/off the tracability.
    """
    global traceEnabled
    traceEnabled = tracable
    if tracable:
        if not logger.handlers:
            logger.addHandler(logging.StreamHandler())
        logger.setLevel(logging.DEBUG)

def setdefaulttimeout(timeout):
    """
    Set the global timeout setting to connect.
    """
    global default_timeout
    default_timeout = timeout

def getdefaulttimeout():
    """
    Return the global timeout setting to connect.
    """
    return default_timeout

def _parse_url(url):
    """
    parse url and the result is tuple of
    (hostname, port, resource path and the flag of secure mode)
    """
    parsed = urlparse(url)
    if parsed.hostname:
        hostname = parsed.hostname
    else:
        raise ValueError("hostname is invalid")
    port = 0
    if parsed.port:
        port = parsed.port

    is_secure = False
    if parsed.scheme == "ws":
        if not port:
            port = 80
    elif parsed.scheme == "wss":
        is_secure = True
        if not port:
            port  = 443
    else:
        raise ValueError("scheme %s is invalid" % parsed.scheme)

    if parsed.path:
        resource = parsed.path
    else:
        resource = "/"

    return (hostname, port, resource, is_secure)

def create_connection(url, timeout=None, **options):
    """
    connect to url and return websocket object.

    Connect to url and return the WebSocket object.
    Passing optional timeout parameter will set the timeout on the socket.
    If no timeout is supplied, the global default timeout setting returned by getdefauttimeout() is used.
    You can customize using 'options'.
    If you set "headers" dict object, you can set your own custom header.

    >>> conn = create_connection("ws://echo.websocket.org/",
    ...     headers={"User-Agent": "MyProgram"})
    """
    websock = WebSocket()
    websock.settimeout(timeout != None and timeout or default_timeout)
    websock.connect(url, **options)
    return websock

_MAX_INTEGER = (1 << 32) -1
_AVAILABLE_KEY_CHARS = range(0x21, 0x2f + 1) + range(0x3a, 0x7e + 1)
_MAX_CHAR_BYTE = (1<<8) -1

# ref. Websocket gets an update, and it breaks stuff.
# http://axod.blogspot.com/2010/06/websocket-gets-update-and-it-breaks.html

def _create_sec_websocket_key():
    spaces_n = random.randint(1, 12)
    max_n = _MAX_INTEGER / spaces_n
    number_n = random.randint(0, max_n)
    product_n = number_n * spaces_n
    key_n = str(product_n)
    for i in range(random.randint(1, 12)):
        c = random.choice(_AVAILABLE_KEY_CHARS)
        pos = random.randint(0, len(key_n))
        key_n = key_n[0:pos] + chr(c) + key_n[pos:]
    for i in range(spaces_n):
        pos = random.randint(1, len(key_n)-1)
        key_n = key_n[0:pos] + " " + key_n[pos:]

    return number_n, key_n

def _create_key3():
    return "".join([chr(random.randint(0, _MAX_CHAR_BYTE)) for i in range(8)])

HEADERS_TO_CHECK = {
    "upgrade": "websocket",
    "connection": "upgrade",
    }

HEADERS_TO_EXIST_FOR_HYBI00 = [
    "sec-websocket-origin",
    "sec-websocket-location",
]

HEADERS_TO_EXIST_FOR_HIXIE75 = [
    "websocket-origin",
    "websocket-location",
]

class _SSLSocketWrapper(object):
    def __init__(self, sock):
        self.ssl = socket.ssl(sock)

    def recv(self, bufsize):
        return self.ssl.read(bufsize)

    def send(self, payload):
        return self.ssl.write(payload)

class WebSocket(object):
    """
    Low level WebSocket interface.
    This class is based on
      The WebSocket protocol draft-hixie-thewebsocketprotocol-76
      http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76

    We can connect to the websocket server and send/recieve data.
    The following example is a echo client.

    >>> import websocket
    >>> ws = websocket.WebSocket()
    >>> ws.connect("ws://echo.websocket.org")
    >>> ws.send("Hello, Server")
    >>> ws.recv()
    'Hello, Server'
    >>> ws.close()
    """
    def __init__(self):
        """
        Initalize WebSocket object.
        """
        self.connected = False
        self.io_sock = self.sock = socket.socket()

    def settimeout(self, timeout):
        """
        Set the timeout to the websocket.
        """
        self.sock.settimeout(timeout)

    def gettimeout(self):
        """
        Get the websocket timeout.
        """
        return self.sock.gettimeout()

    def connect(self, url, **options):
        """
        Connect to url. url is websocket url scheme. ie. ws://host:port/resource
        You can customize using 'options'.
        If you set "headers" dict object, you can set your own custom header.

        >>> ws = WebSocket()
        >>> ws.connect("ws://echo.websocket.org/",
        ...     headers={"User-Agent": "MyProgram"})
        """
        hostname, port, resource, is_secure = _parse_url(url)
        # TODO: we need to support proxy
        self.sock.connect((hostname, port))
        if is_secure:
            self.io_sock = _SSLSocketWrapper(self.sock)
        self._handshake(hostname, port, resource, **options)

    def _handshake(self, host, port, resource, **options):
        sock = self.io_sock
        headers = []
        headers.append("GET %s HTTP/1.1" % resource)
        headers.append("Upgrade: WebSocket")
        headers.append("Connection: Upgrade")
        if port == 80:
            hostport = host
        else:
            hostport = "%s:%d" % (host, port)
        headers.append("Host: %s" % hostport)
        headers.append("Origin: %s" % hostport)

        number_1, key_1 = _create_sec_websocket_key()
        headers.append("Sec-WebSocket-Key1: %s" % key_1)
        number_2, key_2 = _create_sec_websocket_key()
        headers.append("Sec-WebSocket-Key2: %s" % key_2)
        if "header" in options:
            headers.extend(options["header"])

        headers.append("")
        key3 = _create_key3()
        headers.append(key3)

        header_str = "\r\n".join(headers)
        sock.send(header_str)
        if traceEnabled:
            logger.debug( "--- request header ---")
            logger.debug( header_str)
            logger.debug("-----------------------")

        status, resp_headers = self._read_headers()
        if status != 101:
            self.close()
            raise WebSocketException("Handshake Status %d" % status)
        success, secure = self._validate_header(resp_headers)
        if not success:
            self.close()
            raise WebSocketException("Invalid WebSocket Header")

        if secure:
            resp = self._get_resp()
            if not self._validate_resp(number_1, number_2, key3, resp):
                self.close()
                raise WebSocketException("challenge-response error")

        self.connected = True

    def _validate_resp(self, number_1, number_2, key3, resp):
        challenge = struct.pack("!I", number_1)
        challenge += struct.pack("!I", number_2)
        challenge += key3
        digest = hashlib.md5(challenge).digest()

        return  resp == digest

    def _get_resp(self):
        result = self._recv(16)
        if traceEnabled:
            logger.debug("--- challenge response result ---")
            logger.debug(repr(result))
            logger.debug("---------------------------------")

        return result

    def _validate_header(self, headers):
        #TODO: check other headers
        for key, value in HEADERS_TO_CHECK.iteritems():
            v = headers.get(key, None)
            if value != v:
                return False, False

        success = 0
        for key in HEADERS_TO_EXIST_FOR_HYBI00:
            if key in headers:
                success += 1
        if success == len(HEADERS_TO_EXIST_FOR_HYBI00):
            return True, True
        elif success != 0:
            return False, True

        success = 0
        for key in HEADERS_TO_EXIST_FOR_HIXIE75:
            if key in headers:
                success += 1
        if success == len(HEADERS_TO_EXIST_FOR_HIXIE75):
            return True, False

        return False, False

    def _read_headers(self):
        status = None
        headers = {}
        if traceEnabled:
            logger.debug("--- response header ---")

        while True:
            line = self._recv_line()
            if line == "\r\n":
                break
            line = line.strip()
            if traceEnabled:
                logger.debug(line)
            if not status:
                status_info = line.split(" ", 2)
                status = int(status_info[1])
            else:
                kv = line.split(":", 1)
                if len(kv) == 2:
                    key, value = kv
                    headers[key.lower()] = value.strip().lower()
                else:
                    raise WebSocketException("Invalid header")

        if traceEnabled:
            logger.debug("-----------------------")

        return status, headers

    def send(self, payload):
        """
        Send the data as string. payload must be utf-8 string or unicoce.
        """
        if isinstance(payload, unicode):
            payload = payload.encode("utf-8")
        data = "".join(["\x00", payload, "\xff"])
        self.io_sock.send(data)
        if traceEnabled:
            logger.debug("send: " + repr(data))

    def recv(self):
        """
        Reeive utf-8 string data from the server.
        """
        b = self._recv(1)
        if enableTrace:
            logger.debug("recv frame: " + repr(b))
        frame_type = ord(b)
        if frame_type == 0x00:
            bytes = []
            while True:
                b = self._recv(1)
                if b == "\xff":
                    break
                else:
                    bytes.append(b)
            return "".join(bytes)
        elif 0x80 < frame_type < 0xff:
            # which frame type is valid?
            length = self._read_length()
            bytes = self._recv_strict(length)
            return bytes
        elif frame_type == 0xff:
            self._recv(1)
            self._closeInternal()
            return None
        else:
            raise WebSocketException("Invalid frame type")

    def _read_length(self):
        length = 0
        while True:
            b = ord(self._recv(1))
            length = length * (1 << 7) + (b & 0x7f)
            if b < 0x80:
                break

        return length

    def close(self):
        """
        Close Websocket object
        """
        if self.connected:
            try:
                self.io_sock.send("\xff\x00")
                timeout = self.sock.gettimeout()
                self.sock.settimeout(1)
                try:
                    result = self._recv(2)
                    if result != "\xff\x00":
                        logger.error("bad closing Handshake")
                except:
                    pass
                self.sock.settimeout(timeout)
                self.sock.shutdown(socket.SHUT_RDWR)
            except:
                pass
        self._closeInternal()

    def _closeInternal(self):
        self.connected = False
        self.sock.close()
        self.io_sock = self.sock

    def _recv(self, bufsize):
        bytes = self.io_sock.recv(bufsize)
        if not bytes:
            raise ConnectionClosedException()
        return bytes

    def _recv_strict(self, bufsize):
        remaining = bufsize
        bytes = ""
        while remaining:
            bytes += self._recv(remaining)
            remaining = bufsize - len(bytes)

        return bytes

    def _recv_line(self):
        line = []
        while True:
            c = self._recv(1)
            line.append(c)
            if c == "\n":
                break
        return "".join(line)

class WebSocketApp(object):
    """
    Higher level of APIs are provided.
    The interface is like JavaScript WebSocket object.
    """
    def __init__(self, url,
                 on_open = None, on_message = None, on_error = None,
                 on_close = None):
        """
        url: websocket url.
        on_open: callable object which is called at opening websocket.
          this function has one argument. The arugment is this class object.
        on_message: callbale object which is called when recieved data.
         on_message has 2 arguments.
         The 1st arugment is this class object.
         The passing 2nd arugment is utf-8 string which we get from the server.
       on_error: callable object which is called when we get error.
         on_error has 2 arguments.
         The 1st arugment is this class object.
         The passing 2nd arugment is exception object.
       on_close: callable object which is called when closed the connection.
         this function has one argument. The arugment is this class object.
        """
        self.url = url
        self.on_open = on_open
        self.on_message = on_message
        self.on_error = on_error
        self.on_close = on_close
        self.sock = None

    def send(self, data):
        """
        send message. data must be utf-8 string or unicode.
        """
        self.sock.send(data)

    def close(self):
        """
        close websocket connection.
        """
        self.sock.close()

    def run_forever(self):
        """
        run event loop for WebSocket framework.
        This loop is infinite loop and is alive during websocket is available.
        """
        if self.sock:
            raise WebSocketException("socket is already opened")
        try:
            self.sock = WebSocket()
            self.sock.connect(self.url)
            self._run_with_no_err(self.on_open)
            while True:
                data = self.sock.recv()
                if data is None:
                    break
                self._run_with_no_err(self.on_message, data)
        except Exception, e:
            self._run_with_no_err(self.on_error, e)
        finally:
            self.sock.close()
            self._run_with_no_err(self.on_close)
            self.sock = None

    def _run_with_no_err(self, callback, *args):
        if callback:
            try:
                callback(self, *args)
            except Exception, e:
                if logger.isEnabledFor(logging.DEBUG):
                    logger.error(e)

if __name__ == "__main__":
    enableTrace(True)
    ws = create_connection("ws://echo.websocket.org/")
    print "Sending 'Hello, World'..."
    ws.send("Hello, World")
    print "Sent"
    print "Receiving..."
    result =  ws.recv()
    print "Received '%s'" % result
    ws.close()
<commit_msg>- fixed #4 URL parsing is broken<commit_after>"""
websocket - WebSocket client library for Python

Copyright (C) 2010 Hiroki Ohtani(liris)

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

    This library 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
    Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General Public
    License along with this library; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

"""

import socket
from urlparse import urlparse
import random
import struct
import hashlib
import logging

logger = logging.getLogger()

class WebSocketException(Exception):
    pass

class ConnectionClosedException(WebSocketException):
    pass

default_timeout = None
traceEnabled = False

def enableTrace(tracable):
    """
    turn on/off the tracability.
    """
    global traceEnabled
    traceEnabled = tracable
    if tracable:
        if not logger.handlers:
            logger.addHandler(logging.StreamHandler())
        logger.setLevel(logging.DEBUG)

def setdefaulttimeout(timeout):
    """
    Set the global timeout setting to connect.
    """
    global default_timeout
    default_timeout = timeout

def getdefaulttimeout():
    """
    Return the global timeout setting to connect.
    """
    return default_timeout

def _parse_url(url):
    """
    parse url and the result is tuple of
    (hostname, port, resource path and the flag of secure mode)
    """
    if ":" not in url:
        raise ValueError("url is invalid")

    scheme, url = url.split(":", 1)
    url = url.rstrip("/")

    parsed = urlparse(url, scheme="http")
    if parsed.hostname:
        hostname = parsed.hostname
    else:
        raise ValueError("hostname is invalid")
    port = 0
    if parsed.port:
        port = parsed.port

    is_secure = False
    if scheme == "ws":
        if not port:
            port = 80
    elif scheme == "wss":
        is_secure = True
        if not port:
            port  = 443
    else:
        raise ValueError("scheme %s is invalid" % scheme)

    if parsed.path:
        resource = parsed.path
    else:
        resource = "/"

    return (hostname, port, resource, is_secure)

def create_connection(url, timeout=None, **options):
    """
    connect to url and return websocket object.

    Connect to url and return the WebSocket object.
    Passing optional timeout parameter will set the timeout on the socket.
    If no timeout is supplied, the global default timeout setting returned by getdefauttimeout() is used.
    You can customize using 'options'.
    If you set "headers" dict object, you can set your own custom header.

    >>> conn = create_connection("ws://echo.websocket.org/",
    ...     headers={"User-Agent": "MyProgram"})
    """
    websock = WebSocket()
    websock.settimeout(timeout != None and timeout or default_timeout)
    websock.connect(url, **options)
    return websock

_MAX_INTEGER = (1 << 32) -1
_AVAILABLE_KEY_CHARS = range(0x21, 0x2f + 1) + range(0x3a, 0x7e + 1)
_MAX_CHAR_BYTE = (1<<8) -1

# ref. Websocket gets an update, and it breaks stuff.
# http://axod.blogspot.com/2010/06/websocket-gets-update-and-it-breaks.html

def _create_sec_websocket_key():
    spaces_n = random.randint(1, 12)
    max_n = _MAX_INTEGER / spaces_n
    number_n = random.randint(0, max_n)
    product_n = number_n * spaces_n
    key_n = str(product_n)
    for i in range(random.randint(1, 12)):
        c = random.choice(_AVAILABLE_KEY_CHARS)
        pos = random.randint(0, len(key_n))
        key_n = key_n[0:pos] + chr(c) + key_n[pos:]
    for i in range(spaces_n):
        pos = random.randint(1, len(key_n)-1)
        key_n = key_n[0:pos] + " " + key_n[pos:]

    return number_n, key_n

def _create_key3():
    return "".join([chr(random.randint(0, _MAX_CHAR_BYTE)) for i in range(8)])

HEADERS_TO_CHECK = {
    "upgrade": "websocket",
    "connection": "upgrade",
    }

HEADERS_TO_EXIST_FOR_HYBI00 = [
    "sec-websocket-origin",
    "sec-websocket-location",
]

HEADERS_TO_EXIST_FOR_HIXIE75 = [
    "websocket-origin",
    "websocket-location",
]

class _SSLSocketWrapper(object):
    def __init__(self, sock):
        self.ssl = socket.ssl(sock)

    def recv(self, bufsize):
        return self.ssl.read(bufsize)

    def send(self, payload):
        return self.ssl.write(payload)

class WebSocket(object):
    """
    Low level WebSocket interface.
    This class is based on
      The WebSocket protocol draft-hixie-thewebsocketprotocol-76
      http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76

    We can connect to the websocket server and send/recieve data.
    The following example is a echo client.

    >>> import websocket
    >>> ws = websocket.WebSocket()
    >>> ws.connect("ws://echo.websocket.org")
    >>> ws.send("Hello, Server")
    >>> ws.recv()
    'Hello, Server'
    >>> ws.close()
    """
    def __init__(self):
        """
        Initalize WebSocket object.
        """
        self.connected = False
        self.io_sock = self.sock = socket.socket()

    def settimeout(self, timeout):
        """
        Set the timeout to the websocket.
        """
        self.sock.settimeout(timeout)

    def gettimeout(self):
        """
        Get the websocket timeout.
        """
        return self.sock.gettimeout()

    def connect(self, url, **options):
        """
        Connect to url. url is websocket url scheme. ie. ws://host:port/resource
        You can customize using 'options'.
        If you set "headers" dict object, you can set your own custom header.

        >>> ws = WebSocket()
        >>> ws.connect("ws://echo.websocket.org/",
        ...     headers={"User-Agent": "MyProgram"})
        """
        hostname, port, resource, is_secure = _parse_url(url)
        # TODO: we need to support proxy
        self.sock.connect((hostname, port))
        if is_secure:
            self.io_sock = _SSLSocketWrapper(self.sock)
        self._handshake(hostname, port, resource, **options)

    def _handshake(self, host, port, resource, **options):
        sock = self.io_sock
        headers = []
        headers.append("GET %s HTTP/1.1" % resource)
        headers.append("Upgrade: WebSocket")
        headers.append("Connection: Upgrade")
        if port == 80:
            hostport = host
        else:
            hostport = "%s:%d" % (host, port)
        headers.append("Host: %s" % hostport)
        headers.append("Origin: %s" % hostport)

        number_1, key_1 = _create_sec_websocket_key()
        headers.append("Sec-WebSocket-Key1: %s" % key_1)
        number_2, key_2 = _create_sec_websocket_key()
        headers.append("Sec-WebSocket-Key2: %s" % key_2)
        if "header" in options:
            headers.extend(options["header"])

        headers.append("")
        key3 = _create_key3()
        headers.append(key3)

        header_str = "\r\n".join(headers)
        sock.send(header_str)
        if traceEnabled:
            logger.debug( "--- request header ---")
            logger.debug( header_str)
            logger.debug("-----------------------")

        status, resp_headers = self._read_headers()
        if status != 101:
            self.close()
            raise WebSocketException("Handshake Status %d" % status)
        success, secure = self._validate_header(resp_headers)
        if not success:
            self.close()
            raise WebSocketException("Invalid WebSocket Header")

        if secure:
            resp = self._get_resp()
            if not self._validate_resp(number_1, number_2, key3, resp):
                self.close()
                raise WebSocketException("challenge-response error")

        self.connected = True

    def _validate_resp(self, number_1, number_2, key3, resp):
        challenge = struct.pack("!I", number_1)
        challenge += struct.pack("!I", number_2)
        challenge += key3
        digest = hashlib.md5(challenge).digest()

        return  resp == digest

    def _get_resp(self):
        result = self._recv(16)
        if traceEnabled:
            logger.debug("--- challenge response result ---")
            logger.debug(repr(result))
            logger.debug("---------------------------------")

        return result

    def _validate_header(self, headers):
        #TODO: check other headers
        for key, value in HEADERS_TO_CHECK.iteritems():
            v = headers.get(key, None)
            if value != v:
                return False, False

        success = 0
        for key in HEADERS_TO_EXIST_FOR_HYBI00:
            if key in headers:
                success += 1
        if success == len(HEADERS_TO_EXIST_FOR_HYBI00):
            return True, True
        elif success != 0:
            return False, True

        success = 0
        for key in HEADERS_TO_EXIST_FOR_HIXIE75:
            if key in headers:
                success += 1
        if success == len(HEADERS_TO_EXIST_FOR_HIXIE75):
            return True, False

        return False, False

    def _read_headers(self):
        status = None
        headers = {}
        if traceEnabled:
            logger.debug("--- response header ---")

        while True:
            line = self._recv_line()
            if line == "\r\n":
                break
            line = line.strip()
            if traceEnabled:
                logger.debug(line)
            if not status:
                status_info = line.split(" ", 2)
                status = int(status_info[1])
            else:
                kv = line.split(":", 1)
                if len(kv) == 2:
                    key, value = kv
                    headers[key.lower()] = value.strip().lower()
                else:
                    raise WebSocketException("Invalid header")

        if traceEnabled:
            logger.debug("-----------------------")

        return status, headers

    def send(self, payload):
        """
        Send the data as string. payload must be utf-8 string or unicoce.
        """
        if isinstance(payload, unicode):
            payload = payload.encode("utf-8")
        data = "".join(["\x00", payload, "\xff"])
        self.io_sock.send(data)
        if traceEnabled:
            logger.debug("send: " + repr(data))

    def recv(self):
        """
        Reeive utf-8 string data from the server.
        """
        b = self._recv(1)
        if enableTrace:
            logger.debug("recv frame: " + repr(b))
        frame_type = ord(b)
        if frame_type == 0x00:
            bytes = []
            while True:
                b = self._recv(1)
                if b == "\xff":
                    break
                else:
                    bytes.append(b)
            return "".join(bytes)
        elif 0x80 < frame_type < 0xff:
            # which frame type is valid?
            length = self._read_length()
            bytes = self._recv_strict(length)
            return bytes
        elif frame_type == 0xff:
            self._recv(1)
            self._closeInternal()
            return None
        else:
            raise WebSocketException("Invalid frame type")

    def _read_length(self):
        length = 0
        while True:
            b = ord(self._recv(1))
            length = length * (1 << 7) + (b & 0x7f)
            if b < 0x80:
                break

        return length

    def close(self):
        """
        Close Websocket object
        """
        if self.connected:
            try:
                self.io_sock.send("\xff\x00")
                timeout = self.sock.gettimeout()
                self.sock.settimeout(1)
                try:
                    result = self._recv(2)
                    if result != "\xff\x00":
                        logger.error("bad closing Handshake")
                except:
                    pass
                self.sock.settimeout(timeout)
                self.sock.shutdown(socket.SHUT_RDWR)
            except:
                pass
        self._closeInternal()

    def _closeInternal(self):
        self.connected = False
        self.sock.close()
        self.io_sock = self.sock

    def _recv(self, bufsize):
        bytes = self.io_sock.recv(bufsize)
        if not bytes:
            raise ConnectionClosedException()
        return bytes

    def _recv_strict(self, bufsize):
        remaining = bufsize
        bytes = ""
        while remaining:
            bytes += self._recv(remaining)
            remaining = bufsize - len(bytes)

        return bytes

    def _recv_line(self):
        line = []
        while True:
            c = self._recv(1)
            line.append(c)
            if c == "\n":
                break
        return "".join(line)

class WebSocketApp(object):
    """
    Higher level of APIs are provided.
    The interface is like JavaScript WebSocket object.
    """
    def __init__(self, url,
                 on_open = None, on_message = None, on_error = None,
                 on_close = None):
        """
        url: websocket url.
        on_open: callable object which is called at opening websocket.
          this function has one argument. The arugment is this class object.
        on_message: callbale object which is called when recieved data.
         on_message has 2 arguments.
         The 1st arugment is this class object.
         The passing 2nd arugment is utf-8 string which we get from the server.
       on_error: callable object which is called when we get error.
         on_error has 2 arguments.
         The 1st arugment is this class object.
         The passing 2nd arugment is exception object.
       on_close: callable object which is called when closed the connection.
         this function has one argument. The arugment is this class object.
        """
        self.url = url
        self.on_open = on_open
        self.on_message = on_message
        self.on_error = on_error
        self.on_close = on_close
        self.sock = None

    def send(self, data):
        """
        send message. data must be utf-8 string or unicode.
        """
        self.sock.send(data)

    def close(self):
        """
        close websocket connection.
        """
        self.sock.close()

    def run_forever(self):
        """
        run event loop for WebSocket framework.
        This loop is infinite loop and is alive during websocket is available.
        """
        if self.sock:
            raise WebSocketException("socket is already opened")
        try:
            self.sock = WebSocket()
            self.sock.connect(self.url)
            self._run_with_no_err(self.on_open)
            while True:
                data = self.sock.recv()
                if data is None:
                    break
                self._run_with_no_err(self.on_message, data)
        except Exception, e:
            self._run_with_no_err(self.on_error, e)
        finally:
            self.sock.close()
            self._run_with_no_err(self.on_close)
            self.sock = None

    def _run_with_no_err(self, callback, *args):
        if callback:
            try:
                callback(self, *args)
            except Exception, e:
                if logger.isEnabledFor(logging.DEBUG):
                    logger.error(e)

if __name__ == "__main__":
    enableTrace(True)
    ws = create_connection("ws://echo.websocket.org/")
    print "Sending 'Hello, World'..."
    ws.send("Hello, World")
    print "Sent"
    print "Receiving..."
    result =  ws.recv()
    print "Received '%s'" % result
    ws.close()
<|endoftext|>