Add options to break up the captcha image storage with hash-digit subdirectories...
[toast/cookiecaptcha.git] / captcha.py
1 #!/usr/bin/python
2 #
3 # Script to generate distorted text images for a captcha system.
4 #
5 # Copyright (C) 2005 Neil Harris
6 #
7 # This program is free software; you can redistribute it and/or modify
8 # it under the terms of the GNU General Public License as published by
9 # the Free Software Foundation; either version 2 of the License, or
10 # (at your option) any later version.
11 #
12 # This program is distributed in the hope that it will be useful,
13 # but WITHOUT ANY WARRANTY; without even the implied warranty of
14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15 # GNU General Public License for more details.
16 #
17 # You should have received a copy of the GNU General Public License along
18 # with this program; if not, write to the Free Software Foundation, Inc.,
19 # 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
20 # http://www.gnu.org/copyleft/gpl.html
21 #
22 # Further tweaks by Brion Vibber <brion@pobox.com>:
23 # 2006-01-26: Add command-line options for the various parameters
24 # 2007-02-19: Add --dirs param for hash subdirectory splits
25
26 import random
27 import Image
28 import ImageFont
29 import ImageDraw
30 import ImageEnhance
31 import ImageOps
32 import math, string, md5
33 import getopt
34 import os
35 import sys
36
37 # Does X-axis wobbly copy, sandwiched between two rotates
38 def wobbly_copy(src, wob, col, scale, ang):
39         x, y = src.size
40         f = random.uniform(4*scale, 5*scale)
41         p = random.uniform(0, math.pi*2)
42         rr = ang+random.uniform(-30, 30) # vary, but not too much
43         int_d = Image.new('RGB', src.size, 0) # a black rectangle
44         rot = src.rotate(rr, Image.BILINEAR)
45         # Do a cheap bounding-box op here to try to limit work below
46         bbx = rot.getbbox()
47         if bbx == None:
48                 print "whoops"
49                 return src
50         else:
51                 l, t, r, b= bbx
52         # and only do lines with content on
53         for i in range(t, b+1):
54                 # Drop a scan line in
55                 xoff = int(math.sin(p+(i*f/y))*wob)
56                 xoff += int(random.uniform(-wob*0.5, wob*0.5))
57                 int_d.paste(rot.crop((0, i, x, i+1)), (xoff, i))
58         # try to stop blurring from building up
59         int_d = int_d.rotate(-rr, Image.BILINEAR)
60         enh = ImageEnhance.Sharpness(int_d)
61         return enh.enhance(2)
62
63
64 def gen_captcha(text, fontname, fontsize, file_name):
65         """Generate a captcha image"""
66         # white text on a black background
67         bgcolor = 0x0
68         fgcolor = 0xffffff
69         # create a font object 
70         font = ImageFont.truetype(fontname,fontsize)
71         # determine dimensions of the text
72         dim = font.getsize(text)
73         # create a new image significantly larger that the text
74         edge = max(dim[0], dim[1]) + 2*min(dim[0], dim[1])
75         im = Image.new('RGB', (edge, edge), bgcolor)
76         d = ImageDraw.Draw(im)
77         x, y = im.size
78         # add the text to the image
79         d.text((x/2-dim[0]/2, y/2-dim[1]/2), text, font=font, fill=fgcolor)
80         k = 3
81         wob = 0.20*dim[1]/k
82         rot = 45
83         # Apply lots of small stirring operations, rather than a few large ones
84         # in order to get some uniformity of treatment, whilst
85         # maintaining randomness
86         for i in range(k):
87                 im = wobbly_copy(im, wob, bgcolor, i*2+3, rot+0)
88                 im = wobbly_copy(im, wob, bgcolor, i*2+1, rot+45)
89                 im = wobbly_copy(im, wob, bgcolor, i*2+2, rot+90)
90                 rot += 30
91         
92         # now get the bounding box of the nonzero parts of the image
93         bbox = im.getbbox()
94         bord = min(dim[0], dim[1])/4 # a bit of a border
95         im = im.crop((bbox[0]-bord, bbox[1]-bord, bbox[2]+bord, bbox[3]+bord))
96         # and turn into black on white
97         im = ImageOps.invert(im)
98                 
99         # save the image, in format determined from filename
100         im.save(file_name)
101
102 def gen_subdir(basedir, hash, levels):
103         """Generate a subdirectory path out of the first _levels_
104         characters of _hash_, and ensure the directories exist
105         under _basedir_."""
106         subdir = None
107         for i in range(0, levels):
108                 char = hash[i]
109                 if subdir:
110                         subdir = os.path.join(subdir, char)
111                 else:
112                         subdir = char
113                 fulldir = os.path.join(basedir, subdir)
114                 if not os.path.exists(fulldir):
115                         os.mkdir(fulldir)
116         return subdir
117         
118 if __name__ == '__main__':
119         """This grabs random words from the dictionary 'words' (one
120         word per line) and generates a captcha image for each one,
121         with a keyed salted hash of the correct answer in the filename.
122         
123         To check a reply, hash it in the same way with the same salt and
124         secret key, then compare with the hash value given.
125         """
126         font = "VeraBd.ttf"
127         wordlist = "awordlist.txt"
128         key = "CHANGE_THIS_SECRET!"
129         output = "."
130         count = 20
131         fill = 0
132         dirs = 0
133         verbose = False
134         
135         opts, args = getopt.getopt(sys.argv[1:], "", ["font=", "wordlist=", "key=", "output=", "count=", "fill=", "dirs=", "verbose"])
136         for o, a in opts:
137                 if o == "--font":
138                         font = a
139                 if o == "--wordlist":
140                         wordlist = a
141                 if o == "--key":
142                         key = a
143                 if o == "--output":
144                         output = a
145                 if o == "--count":
146                         count = int(a)
147                 if o == "--fill":
148                         fill = int(a)
149                 if o == "--dirs":
150                         dirs = int(a)
151                 if o == "--verbose":
152                         verbose = True
153         
154         if fill:
155                 # Option processing order is not guaranteed, so count the output
156                 # files after...
157                 count = max(0, fill - len(os.listdir(output)))
158         
159         words = [string.lower(x.strip()) for x in open(wordlist).readlines()]
160         words = [x for x in words
161                 if len(x) <= 5 and len(x) >= 4 and x[0] != "f"
162                 and x[0] != x[1] and x[-1] != x[-2]
163                 and (not "'" in x)]
164         for i in range(count):
165                 word1 = words[random.randint(0,len(words)-1)]
166                 word2 = words[random.randint(0,len(words)-1)]
167                 word = word1+word2
168                 salt = "%08x" % random.randrange(2**32)
169                 # 64 bits of hash is plenty for this purpose
170                 hash = md5.new(key+salt+word+key+salt).hexdigest()[:16]
171                 filename = "image_%s_%s.png" % (salt, hash)
172                 if dirs:
173                         subdir = gen_subdir(output, hash, dirs)
174                         filename = os.path.join(subdir, filename)
175                 if verbose:
176                         print filename
177                 gen_captcha(word, font, 40, os.path.join(output, filename))