1- from io import BytesIO
1+ #!/usr/bin/env python3
2+ # -*- coding: utf-8 -*-
23import argparse
3- import os
4- import requests
4+ import logging
5+ import re
6+ import secrets
57import subprocess
68import sys
79import zipfile
10+ from base64 import b64encode , b64decode
11+ from hashlib import sha256
12+ from io import BytesIO
813from pathlib import Path
914
10- ##
11- ## Este módulo deve ser executado na raiz do projeto
12- ##
15+ import requests
16+ from kazoo .client import KazooClient
17+
18+ #
19+ # Este módulo deve ser executado na raiz do projeto
20+ #
21+
22+ logging .basicConfig ()
23+
24+ SECURITY_FILE_TEMPLATE = """
25+ {
26+ "authentication":{
27+ "blockUnknown": true,
28+ "class":"solr.BasicAuthPlugin",
29+ "credentials":{"%s":"%s %s"},
30+ "forwardCredentials": false,
31+ "realm": "Solr Login"
32+ },
33+ "authorization":{
34+ "class":"solr.RuleBasedAuthorizationPlugin",
35+ "permissions":[{"name":"security-edit", "role":"admin"}],
36+ "user-role":{"%s":"admin"}
37+ }
38+ }
39+ """
40+
41+ URL_PATTERN = 'https?://(([a-zA-Z0-9]+):([a-zA-Z0-9]+)@)?([a-zA-Z0-9.-]+)(:[0-9]{4})?'
42+
43+
44+ def solr_hash_password (password : str , salt : str = None ):
45+ """
46+ Generates a password and salt to be used in Basic Auth Solr
47+
48+ password: clean text password string
49+ salt (optional): base64 salt string
50+ returns: sha256 hash of password and salt (both base64 strings)
51+ """
52+ m = sha256 ()
53+ if salt is None :
54+ salt = secrets .token_bytes (32 )
55+ else :
56+ salt = b64decode (salt )
57+ m .update (salt + password .encode ('utf-8' ))
58+ digest = m .digest ()
1359
14- class SolrClient :
60+ m = sha256 ()
61+ m .update (digest )
62+ digest = m .digest ()
63+
64+ cypher = b64encode (digest ).decode ('utf-8' )
65+ salt = b64encode (salt ).decode ('utf-8' )
66+ return cypher , salt
67+
68+
69+ def create_security_file (username , password ):
70+ print ("Creating security.json file..." )
71+ with open ("security.json" , "w" ) as f :
72+ cypher , salt = solr_hash_password (password )
73+ f .write (SECURITY_FILE_TEMPLATE % (username , cypher , salt , username ))
74+ print ("file created!" )
75+
76+
77+ def upload_security_file (zk_host ):
78+ zk_port = 9983 # embedded ZK port
79+ print (f"Uploading security file to Solr, ZK server={ zk_host } { zk_port } )
80+ try :
81+ with open ('security.json' , 'r' ) as f :
82+ data = f .read ()
83+ zk = KazooClient (hosts = f"{ zk_host } { zk_port } )
84+ zk .start ()
85+ print ("Uploading security.json file..." )
86+ if zk .exists ('/security.json' ):
87+ zk .set ("/security.json" , str .encode (data ))
88+ else :
89+ zk .create ("/security.json" , str .encode (data ))
90+ data , stat = zk .get ('/security.json' )
91+ print ("file uploaded!" )
92+ print (data .decode ('utf-8' ))
93+ zk .stop ()
94+ except Exception as e :
95+ print (e )
96+ sys .exit (- 1 )
1597
98+
99+ class SolrClient :
16100 LIST_CONFIGSETS = "{}/solr/admin/configs?action=LIST&omitHeader=true&wt=json"
17101 UPLOAD_CONFIGSET = "{}/solr/admin/configs?action=UPLOAD&name={}&wt=json"
18102 LIST_COLLECTIONS = "{}/solr/admin/collections?action=LIST&wt=json"
@@ -129,7 +213,7 @@ def create_collection(self, collection_name, shards=1, replication_factor=1, max
129213 print (res .content )
130214 return False
131215 return True
132-
216+
133217 def delete_collection (self , collection_name ):
134218 if collection_name == '*' :
135219 collections = self .list_collections ()
@@ -160,6 +244,22 @@ def delete_index_data(self, collection_name):
160244 print ("Num docs: %s" % num_docs )
161245
162246
247+ def setup_embedded_zk (solr_url ):
248+ match = re .match (URL_PATTERN , solr_url )
249+ if match :
250+ _ , solr_user , solr_pwd , solr_host , solr_port = match .groups ()
251+
252+ if solr_user and solr_pwd and solr_host :
253+ create_security_file (solr_user , solr_pwd )
254+ upload_security_file (solr_host )
255+ else :
256+ print (f"Missing Solr's username, password, and host: { solr_user } { solr_pwd } { solr_host } )
257+ sys .exit (- 1 )
258+ else :
259+ print (f"Solr URL path doesn't match the required format: { solr_url } )
260+ sys .exit (- 1 )
261+
262+
163263if __name__ == '__main__' :
164264
165265 parser = argparse .ArgumentParser (description = 'Cria uma collection no Solr' )
@@ -178,17 +278,27 @@ def delete_index_data(self, collection_name):
178278 parser .add_argument ('-ms' , type = int , dest = 'max_shards_per_node' , nargs = '?' ,
179279 help = 'Max shards per node (default=1)' , default = 1 )
180280
281+ parser .add_argument ("--embedded_zk" , default = False , action = "store_true" ,
282+ help = "Embedded ZooKeeper" )
283+
181284 try :
182285 args = parser .parse_args ()
183286 except IOError as msg :
184287 parser .error (str (msg ))
185288 sys .exit (- 1 )
186289
187290 url = args .url .pop ()
188- collection = args .collection .pop ()
189291
292+ if args .embedded_zk :
293+ print ("Setup embedded ZooKeeper..." )
294+ setup_embedded_zk (url )
295+
296+ collection = args .collection .pop ()
190297 client = SolrClient (url = url )
191298
299+ ## Add --force to force upload security.json, configset upload and collection recreation
300+ ## it will clean the solr server before proceeding
301+ ## Add --clean option to clean uploadconfig and collection
192302 if not client .exists_collection (collection ):
193303 print ("Collection '%s' doesn't exists. Creating a new one..." % collection )
194304 created = client .create_collection (collection ,
@@ -200,6 +310,7 @@ def delete_index_data(self, collection_name):
200310 else :
201311 print ("Collection '%s' exists." % collection )
202312
313+ ## Add --disable-index to disable auto index
203314 num_docs = client .get_num_docs (collection )
204315 if num_docs == 0 :
205316 print ("Performing a full reindex of '%s' collection..." % collection )
0 commit comments