import urllib.request
import urllib.error
import struct
import re
import webbrowser
import traceback
import sys
import os
import hashlib
import base64
import time

NON_B = 'a c d e g gif h hr k m o p r s t u v w wg i ic cm y 3 adv an cgl ck co fa fit int jp lit mu n po sci soc sp tg toy trv tv vp x'.split(' ')
SAVE_PATH = 'found'
LOG_FILE = 'log.txt'
SHORT_LENGTH = 100
ENDS_LENGTH = 4
MIN_JPEG_FF = 3
MAX_JPG_WO_FF = 1000
MAX_NON_ASCII = 0.01
MIN_FRAC_JPEG = 0.9

scanned = set(fag)

def content_found(cat, info, board, post, uri, data, log):
    """Called when hidden content is found."""
    if not cat in ['nothing', 'unsupported_type']:
        if cat in ['4chan_sounds', 'WinRAR', 'ChanGrouper', 'wetherbed', 'HTML', 'text', 'JPEG', 'MPEG_audio', 'unrecognized']:
            name, new = save_data(cat, data, uri.split('.')[-1])
            if new:
                print(post, uri, cat, info, '->', name, file=log)
                webbrowser.open(post_uri(board, post))
            else:
                print(post, uri, cat, info, ':', name, file=log)
        else:
            print(post, uri, cat, info, file=log)

def page_uri(board, page=0):
    return 'http://boards.4chan.org/' + board + '/' + ('' if page == 0 else str(page))

def post_uri(board, post):
    return 'http://sys.4chan.org/' + board + '/imgboard.php?res=' + post

def save_data(cat, data, ext):
    """
    Save file in appropriate directory.
    Return name, whether file was new.
    """
    if not os.path.exists(SAVE_PATH):
        os.mkdir(SAVE_PATH)
    if not os.path.exists(SAVE_PATH + '/' + cat):
        os.mkdir(SAVE_PATH + '/' + cat)
    hash = base64.b64encode(hashlib.md5(data).digest(),b'-_')[:-2].decode('ASCII')
    name = SAVE_PATH + '/' + cat + '/' + hash + '.' + ext
    if not os.path.exists(name):
        with open(name, 'wb') as f:
            f.write(data)
        return name, True
    else:
        return name, False

def color_table_length(c):
    """Read the flags in a GIF file to get the length of the color table."""
    if c & 0x80:
        return 3 << ((c & 0x07) + 1)
    else:
        return 0

def end_of_blocks(data, start):
    """Find the end of a sequence of GIF data sub-blocks."""
    i = start
    while i < len(data):
        if data[i] == 0:
            return i + 1
        else:
            i += data[i] + 1
    return len(data)

def gif_end(data):
    """Find the end of GIF data in a file."""
    if len(data) <= 10:
        return len(data)
    i = color_table_length(data[10]) + 13
    while i < len(data):
        if data[i] == 0x2c:
            i += 9
            if i >= len(data):
                return len(data)
            i += color_table_length(data[i]) + 2
            i = end_of_blocks(data, i)
        elif data[i] == 0x21:
            i = end_of_blocks(data, i+2)
        elif data[i] == 0x3b:
            return i + 1
        else:
            return i
    return min(i, len(data))

def jpg_end(data):
    """Find the end of JPEG data in a file."""
    i = 2
    while i + 2 <= len(data):
        if data[i] == 0xff and data[i+1] >= 0xc0:
            if data[i+1] == 0xd9:
                return i + 2
            elif data[i+1] == 0xd8:
                return i
            elif 0xD0 <= data[i+1] <= 0xD7:
                i += 2
            else:
                if i + 4 > len(data):
                    return len(data)
                i += struct.unpack_from('>H', data, i+2)[0] + 2
        else:
            i += 1
    return min(i, len(data))

def png_end(data):
    """Find the end of PNG data in a file."""
    i = 8
    while i + 8 < len(data):
        length, chunk_type = struct.unpack_from('>L4s', data, i)
        if chunk_type == b'IEND':
            return min(i + 12, len(data))
        i += length + 12
    return min(i, len(data))

ENDF = {'gif': gif_end, 'jpg': jpg_end, 'png': png_end}
TERM = {'gif': b';', 'jpg': b'\xff\xd9', 'png': b'\0\0\0\0IEND\xaeB`\x82'}

MPEG_BITRATE = [
16*[-1],
[-1,8,16,24,32,40,48,56,64,80,96,112,128,144,160,-1],
[-1,8,16,24,32,40,48,56,64,80,96,112,128,144,160,-1],
[-1,32,48,56,64,80,96,112,128,144,160,176,192,224,256,-1],
16*[-1],
[-1,32,40,48,56,64,80,96,112,128,160,192,224,256,320,-1],
[-1,32,48,56,64,80,96,112,128,160,192,224,256,320,384,-1],
[-1,32,64,96,128,160,192,224,256,288,320,352,384,416,448,-1]]

MPEG_SAMPLINGRATE = [
[11025,12000,8000,-1],
4*[-1],
[22050,24000,16000,-1],
[44100,48000,32000,-1]]

def cat_hidden(data, ext):
    """
    Determines type of appended data.
    Returns type name, other info to print.
    """
    if not ext in ENDF:
        return 'unsupported_type', ''

    if data[-4:] == b'\xa1\xb2\xc3\xd4':
        return 'FotoStation', ''
    if data[-12:-8] == b'phtv':
        return 'phtv', ''
    if data[-12:-8] == b'AXS!' or data[-12:-8] == b'AXS*':
        return 'AFCP', ''
    if data[-8:] == b'cbipcbbl':
        return 'PhotoMechanic', ''
    if data[-9:] == b'</pImgDB>':
        return 'pImgDB', ''
    if data.find(b'Motorola', 0, 478) != -1:
        return 'Droid', ''

    end = ENDF[ext](data)

    if end == len(data):
        return 'nothing', ''
    if all(c in [9, 10, 13, 32] for c in data[end:]):
        return 'whitespace', ''
    if end == len(data) - 1:
        return 'single_byte', ''
    if data[end:] == bytes((len(data) - end) * [data[-1]]):
        return 'repeating_byte', ''
    if data[end:end+2] == b'\x10\x88' and not any(data[end+2:]):
        return '10_88_nulls', ''
    if data[end:end+29] == b'This file has been marked by ':
        return 'cagedgirl', ''
    if data[end:end+16] == b'\x80\x3F\xE0\x50\x38\x24\x16\x0D\x07\x84\x42\x61\x50\xB8\x64\x36':
        return '80_3F_E0_50', ''
    if data[end:end+11] == b'SplashPhoto':
        return 'SplashPhoto', ''

    if data[end] == 91 and data.find(b'OggS', end, end+9):
        return '4chan_sounds', ''
    if data.find(b'ar!\x1A\x07\x00') != -1:
        return 'WinRAR', ''
    if data.find(b'<!--Path Info: /i4M/file_view/') != -1:
        return 'wetherbed', ''
    if data.find(b'eWallpapers.eu', end) != -1 or data.find(b'ewallpapers.eu', end) != -1:
        return 'eWallpapers', ''
    if data.find(b'Box The Clown Image Hosting', end) != -1:
        return 'Box_The_Clown', ''

    if len(data) - end >= 8:
        i = struct.unpack_from('>q', data, -8)[0]
        if end <= i <= len(data) - 8:
            i += struct.unpack_from('>I', data, i)[0] + 4
            while i < len(data) - 8:
                i += struct.unpack_from('>I', data, i)[0] + 4
                if i > len(data) - 8:
                    break
                i += struct.unpack_from('>Q', data, i)[0] + 8
            if i == len(data) - 8:
                return 'ChanGrouper', ''

    n = 0
    nmax = MAX_NON_ASCII * (len(data) - end)
    for c in data[end:]:
        if not (0x20 <= c <= 0x7e or c in [9,10,13]):
            n += 1
            if n > nmax:
                break
    else:
        if data[end:end+6] == b'\xcc\xbb\r\n\r\n':
            return 'text_metadata', ''
        elif data.find(b'</html>', end) != -1 or data.find(b'</body>', end) != -1 or data.find(b'<br>', end) != -1 or data.find(b'<br />', end) != -1 :
            return 'HTML', ''
        elif len(data) - end <= SHORT_LENGTH:
            return 'text_short', str(len(data) - end) + ' ' + hex(end) + ' ' + ascii(data[end:])
        else:
            return 'text', str(len(data) - end) + ' ' + hex(end) + ' ' + summary(data[end:])

    i = data.find(b'\xff\xd8\xff', end)
    if i != -1:
        j = i + jpg_end(data[i:])
        if data[j-2:j] == b'\xff\xd9' and (j-i) / (len(data)-end) >= MIN_FRAC_JPEG:
            return 'JPEG', ''
        elif len(data) - j < 2 and (len(data)-i) / (len(data)-end) >= MIN_FRAC_JPEG:
            return 'JPEG_start_frag', ''

    n = 0
    for i in range(end, len(data)-1):
        if data[i] == 0xff:
            if not (data[i+1] == 0 or 0xd0 <= data[i+1] <= 0xd7 or (i >= len(data) - 3 and data[i+1] == 0xd9)):
                break
            n += 1
    else:
        if n >= MIN_JPEG_FF:
            i = data.find(b'\xff', end)
            j = data.rfind(b'\xff', end)
            if (j-i+MAX_JPG_WO_FF) / (len(data)-end) >= MIN_FRAC_JPEG:
                if j + 1 < len(data) and data[j+1] == 0xd9:
                    return 'JPEG_end_frag', ''
                else:
                    return 'JPEG_mid_frag', ''

    frames = 0
    for i in range(end, len(data)-3):
        if data[i] == 0xff:
            if data[i+1] & 0xE0 == 0xE0:
                break
    while i + 4 <= len(data) and data[i] == 0xFF and data[i+1] & 0xE0 == 0xE0:
        bitrate = MPEG_BITRATE[(data[i+1] & 0x0E) >> 1][(data[i+2] & 0xF0) >> 4]
        samplingrate = MPEG_SAMPLINGRATE[(data[i+1] & 0x18) >> 3][(data[i+2] & 0x0C) >> 2]
        if bitrate < 0 or samplingrate < 0:
            break
        padding = (data[i+2] & 0x02) >> 1
        if data[i+1] & 0x06 == 0x06:
            n = 4 * (int(12000*bitrate/samplingrate) + padding);
        else:
            n = int(144000*bitrate/samplingrate) + padding;
        if i + n <= len(data):
            i += n
            frames += 1
        else:
            break
    if frames >= 4:
        return 'MPEG_audio', ''

    if len(data) - end <= SHORT_LENGTH:
        return 'unrecognized_short', str(len(data) - end) + ' ' + hex(end) + ' ' + ascii(data[end:])
    else:
        return 'unrecognized', str(len(data) - end) + ' ' + hex(end) + ' ' + summary(data[end:])

def summary(data):
    if data[:ENDS_LENGTH] == ENDS_LENGTH * b'\x00':
        i = 0
        while i < len(data) - 2*ENDS_LENGTH and data[i] == 0:
            i += 1
        return str(i) + " null " + ascii(data[i:i+ENDS_LENGTH]) + ' ... ' + ascii(data[-ENDS_LENGTH:])
    else:
        return ascii(data[:ENDS_LENGTH]) + ' ... ' + ascii(data[-ENDS_LENGTH:])

def scan_uri(uri):
    """
    Quickly determine whether a URI is likely to contain hidden content.
    Returns type name, additional info, and response from URI if fully downloaded.
    """
    ext = uri.split('.')[-1]
    if not ext in ENDF:
        return 'unsupported_type', '', b''
    req = urllib.request.Request(uri)
    req.add_header('Range', 'bytes=-' + str(len(TERM[ext])))
    response = urllib.request.urlopen(req)
    if response.read() == TERM[ext]:
        return 'nothing', '', b''
    else:
        rfull = urllib.request.urlopen(uri)
        data = rfull.read()
        cat, info = cat_hidden(data, ext)
        sys.stdout.flush()
        return cat, info, data

def scan_page(board, page, log):
    """Check all the images on a page for hidden content."""
    try:
        page_content = bytes.decode(urllib.request.urlopen(page).read())
    except urllib.error.URLError:
        return
    matches = re.findall(r'<a href="([^"]+)" target=_blank><img[^>]*></a><a name="0"></a>\n<input type=checkbox name="(\d+)"', page_content)
    matches = [(x[1], x[0]) for x in matches]
    matches += re.findall(r'>(\d+)</a></span><br>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;<span class="filesize">File<a href="([^"]+)"', page_content)
    for match in matches:
        post, uri = match
        try:
            if not (board, post) in scanned:
                try:
                    cat, info, data = scan_uri(uri)
                except urllib.error.URLError:
                    continue
                scanned.add((board, post))
                content_found(cat, info, board, post, uri, data, log)
        except KeyboardInterrupt:
            raise
        except:
            print('Error on', post, uri, file=log)
            traceback.print_exc(file=log)

def scan_file(filename):
    """Scan file.  Return type, info, file contents."""
    ext = filename.split('.')[-1]
    with open(filename, 'rb') as f:
        data = f.read()
    cat, info = cat_hidden(data, ext)
    return cat, info, data

class LogStream:
    def __init__(self):
        self.f = open(LOG_FILE, 'a')
    def __del__(self):
        self.f.close()
    def write(self, text):
        sys.stdout.write(text)
        sys.stdout.flush()
        self.f.write(text)
        self.f.flush()
    def flush(self):
        pass

def scan_boards():
    """Scan 4chan boards for hidden content."""
    log = LogStream()
    print('\nScan started', time.asctime(time.gmtime()), file=log)
    while True:
        for board in NON_B:
            try:
                scan_page('b', page_uri('b'), log)
                scan_page(board, page_uri(board), log)
            except KeyboardInterrupt:
                print('Scan ended', time.asctime(time.gmtime()), file=log)
                exit()
            except:
                traceback.print_exc(file=log)

def add_from_dir(path):
    """Classify files in directory and copy to save folder."""
    for fn in os.listdir(path):
        fullname = path + '/' + fn
        if os.path.isfile(fullname):
            cat, info, data = scan_file(fullname)
            if cat != 'unsupported_type':
                name, new = save_data(cat, data, fn.split('.')[-1])
                print(fn, '->' if new else ':', name)
                sys.stdout.flush()

def redo():
    """Redo classification of saved files."""
    for sub in os.listdir(SAVE_PATH):
        for fn in os.listdir(SAVE_PATH + '/' + sub):
            fullname = SAVE_PATH + '/' + sub + '/' + fn
            cat, info, data = scan_file(fullname)
            if cat != 'unsupported_type' and cat != sub:
                if not os.path.exists(SAVE_PATH + '/' + cat):
                    os.mkdir(SAVE_PATH + '/' + cat)
                newname = SAVE_PATH + '/' + cat + '/' + fn
                if os.path.exists(newname):
                    os.remove(fullname)
                    print(fullname, ':', cat)
                else:
                    os.rename(fullname, newname)
                    print(fullname, '->', cat)
                sys.stdout.flush()

if __name__ == '__main__':
    if len(sys.argv) >= 3 and sys.argv[1] == 'file':
        cat, info, data = scan_file(sys.argv[2])
        print(sys.argv[2], cat, info)
        sys.stdout.flush()
    elif len(sys.argv) >= 3 and sys.argv[1] == 'dir':
        for fn in os.listdir(sys.argv[2]):
            cat, info, data = scan_file(sys.argv[2] + '/' + fn)
            if cat != 'unsupported_type':
                print(fn, cat, info)
                sys.stdout.flush()
    elif len(sys.argv) >= 2 and sys.argv[1] == 'scan':
        scan_boards()
    elif len(sys.argv) >= 3 and sys.argv[1] == 'add':
        add_from_dir(sys.argv[2])
    elif len(sys.argv) >= 2 and sys.argv[1] == 'redo':
        redo()
    else:
        print("""
Commands:
    scan                    Scan 4chan boards for data appended to images.
    file <path/filename>    Scan file; print classification.
    dir <pathname>          Scan files in directory; print classification.
    add <pathname>          Classify files in directory; copy to save folder.
    redo                    Reclassify files in save folder.
""")