# coding: utf-8
# captcha.py
# http://www.site-digger.com
# hello@site-digger.com
# Identify captcha on http://www.sn.10086.cn/

import Image

class Captcha:
    def __init__(self):
        self.imglib = {}
        self._loadlib()
    
    def _loadlib(self):
        """Load characteristic image lib"""
        
        import os
        if not os.path.exists('imglib'):
            print 'Can not find imglib dir.'
            return
        
        for i in range(10):
            self.imglib[i] = []
            img = Image.open('imglib/%d.jpg' % i).convert('1')
            width, height = img.size
            for w in range(width):
                # store all pixels in a column
                column = []
                for h in range(height):
                    column.append(img.getpixel((w, h)))
                self.imglib[i].append(column)
        
    def _cmpmatrix(self, listA, listB):
        """Return the count of difference between two list"""
        
        if len(listA) != len(listB): return
        
        num = 0
        for i, column in enumerate(listA):
            if len(column) != len(listB[i]): return
            for j, pixel in enumerate(column):
                if pixel != listB[i][j]:
                    num += 1
        return num
    
    def _whichnum(self, piexls_matrix):
        """Identify single number"""

        minnum = None
        index = 0
        for i in range(10):
            ret = self._cmpmatrix(self.imglib.get(i, []), piexls_matrix)
            if ret!= None:
                if minnum == None or minnum > ret:
                    minnum = ret
                    index = i
  
        if minnum != None:
            return str(index)
        else:
            return '?'
    
    def identify(self, filepath=None, fileobj=None):
        """Identify captcha"""

        if filepath:
            img = Image.open(filepath)
        elif fileobj:
            img = Image.open(fileobj)
        else:
            print 'Invalid input.'
            return
        
        img = img.convert('1')
        
        width, height = img.size
        
        w = 0
        number = ''
        while w < width:
            column = []
            for h in range(height):
                column.append(img.getpixel((w, h)))

            # begining of a number
            if sum(column)/height < 245:
                piexls_matrix = []
                for i in range(8):
                    piexls_column = []
                    for j in range(20):
                        piexls_column.append(img.getpixel((w + i, j)))
                    piexls_matrix.append(piexls_column)
                    
                number += self._whichnum(piexls_matrix)
                w = w + 10
            else:
                w = w + 1

        return number
       
if __name__ == '__main__':
    """Test performance of Captcha Class"""
    captcha = Captcha()

    try:
        import urllib2
        response = urllib2.urlopen('https://sn.ac.10086.cn/SSO/servlet/CreateImage')
        open('code.jpg', 'wb').write(response.read())
        Image.open('code.jpg').show()
        print captcha.identify('code.jpg')
    except Exception, e:
        print 'Download captcha fail.', e