forked from MISP/misp-modules
-
Notifications
You must be signed in to change notification settings - Fork 0
/
ocr.py
executable file
·112 lines (91 loc) · 3.08 KB
/
ocr.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import sys
import json
import base64
from io import BytesIO
import logging
log = logging.getLogger('ocr')
log.setLevel(logging.DEBUG)
ch = logging.StreamHandler(sys.stdout)
ch.setLevel(logging.DEBUG)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
ch.setFormatter(formatter)
log.addHandler(ch)
misperrors = {'error': 'Error'}
userConfig = {}
inputSource = ['file']
moduleinfo = {'version': '0.2', 'author': 'Alexandre Dulaunoy',
'description': 'Optical Character Recognition (OCR) module for MISP',
'module-type': ['import']}
moduleconfig = []
def handler(q=False):
# try to import modules and return errors if module not found
try:
from PIL import Image
except ImportError:
misperrors['error'] = "Please pip(3) install pillow"
return misperrors
try:
# Official ImageMagick module
from wand.image import Image as WImage
except ImportError:
misperrors['error'] = "Please pip(3) install wand"
return misperrors
try:
from pytesseract import image_to_string
except ImportError:
misperrors['error'] = "Please pip(3) install pytesseract"
return misperrors
if q is False:
return False
r = {'results': []}
request = json.loads(q)
document = base64.b64decode(request["data"])
document = WImage(blob=document)
if document.format == 'PDF':
with document as pdf:
# Get number of pages
pages = len(pdf.sequence)
log.debug("PDF with {} page(s) detected".format(pages))
# Create new image object where the height will be the number of pages. With huge PDFs this will overflow, break, consume silly memory etc…
img = WImage(width=pdf.width, height=pdf.height * pages)
# Cycle through pages and stitch it together to one big file
for p in range(pages):
log.debug("Stitching page {}".format(p + 1))
image = img.composite(pdf.sequence[p], top=pdf.height * p, left=0)
# Create a png blob
image = img.make_blob('png')
log.debug("Final image size is {}x{}".format(pdf.width, pdf.height * (p + 1)))
else:
image = base64.b64decode(request["data"])
image_file = BytesIO(image)
image_file.seek(0)
try:
im = Image.open(image_file)
except IOError:
misperrors['error'] = "Corrupt or not an image file."
return misperrors
ocrized = image_to_string(im)
freetext = {}
freetext['values'] = ocrized
freetext['types'] = ['freetext']
r['results'].append(freetext)
return r
def introspection():
modulesetup = {}
try:
userConfig
modulesetup['userConfig'] = userConfig
except NameError:
pass
try:
inputSource
modulesetup['inputSource'] = inputSource
except NameError:
pass
return modulesetup
def version():
moduleinfo['config'] = moduleconfig
return moduleinfo
if __name__ == '__main__':
x = open('test.json', 'r')
handler(q=x.read())