Gitingest

Include files under: 50kB

Private Repository NEW

Summary

Repository: pynow/elasticsearch
Files analyzed: 4

Estimated tokens: 4.0k

Download

Directory Structure

Directory structure:

└── pynow-elasticsearch/

    ├── README.md

    ├── es2json.py

    ├── fdns_download.py

    └── json2es.py

Files Content

================================================
FILE: README.md
================================================
#elasticsearch Python脚本

---

###es2json.py    
用于将数据从elasticsearch导出为json    
支持导出过程中将json文件压缩为gz格式

---   

###json2es.py    
用于将json格式文件的输入导入elasticsearch    
支持json文件的gz压缩包直接导入    
支持自定义mapping（在脚本所在目录加入mapping.json）

---

###fdns_download.py    
下载scans.io下的最新dnsrecords_all数据，自动解压入库    
网络数据连接：https://scans.io/study/sonar.fdns



================================================
FILE: es2json.py
================================================

#coding: utf-8
import sys
import json
import gzip
import getopt
import threading
import Queue
from datetime import datetime
from elasticsearch import Elasticsearch, helpers

class logRecode(object):
    def __init__(self):
        self.begin_time = datetime.now()
        self.end_time = datetime.now()
        self.total = 0
        self.success = 0 
        self.failed = 0   

class ElasticExportIndex(logRecode):
    
    def __init__(self, es_host, index_name, save_name,
                  b_gzip=False, count=False, scroll_size=1000,queue_size=10000):  
        logRecode.__init__(self)  
        self.elastic = Elasticsearch(es_host)
        self.index_name = index_name
        self.save_name = save_name
        self.b_gzip = b_gzip
        self.count = int(count) if count else False # select size
        self.res_count = 0
        self.scroll_size = scroll_size                
        self.queue = Queue.Queue(queue_size)
        self.fquit = object()          
        self.start()                        
        
    def read_elastic(self):                                            
        body = {'query': {'match_all':{}}}            
        res = self.elastic.search(index=self.index_name, 
                        body=body, 
                        scroll='3m',
                        search_type='scan',
                        size=self.scroll_size,
                        request_timeout=60)             
        scroll_size = res['hits']['total']              
        self.total = scroll_size
        # print self.total
        
        while (scroll_size > 0):
            
            try:                   
                scroll_id = res['_scroll_id']
                res = self.elastic.scroll(scroll_id=scroll_id, scroll='3m')                  
                for record in res['hits']['hits']:
                    try:                                                                                                       
                        self.queue.put(record['_source'])
                    except Exception as e:                  
                        print e                                                                            
                scroll_size = len(res['hits']['hits'])
                self.res_count += scroll_size # count check
                if self.count and (self.count <= self.res_count): break
                # print self.res_count,self.count
                print('[+] fetch {}...'.format(self.res_count))
            except Exception as e:                              
                pass                              
    
    def write_file(self):       
        try:
            if self.b_gzip:
                fp = gzip.open(self.save_name + '.gz', 'wb')
            else:
                fp = open(self.save_name, "wb")                
            while True:                
                record = self.queue.get()
                if record == self.fquit:
                    self.queue.put(self.fquit)
                    break
                try:                 
                    fp.write(json.dumps(record) + '\n')
                    self.success += 1
                except Exception, e:
                    self.failed += 1
                    print e
                self.queue.task_done()           
        finally:
            fp.close()
        self.end_time = datetime.now()
    
    def __str__(self):
        return json.dumps({
            "begin_time": str(self.begin_time),
            "end_time:": str(self.end_time),
            "total:": self.total,
            "success:": self.success,
            "failed:":  self.failed}, indent=4)      
    
    def start(self):
        if not self.elastic.indices.exists(self.index_name):
            print  "Error: Elasticsearch index %s not exists." % self.index_name
            sys.exit(-1)
            
        read_thread  = threading.Thread(target=self.read_elastic)
        read_thread.start()           
         
        write_thread = threading.Thread(target=self.write_file)
        write_thread.start()
        
        read_thread.join()
        self.queue.join()
        self.queue.put(self.fquit)
        write_thread.join()                      
    
def usage():      
    print "Example:"
    print "%s -e 192.168.1.121:9200 -i indexname -f result.json -c 10000 -g " %sys.argv[0]    
    sys.exit(1)
    
def main():    
    if len(sys.argv[1:]) < 1:
        usage()
    try: 
        opts, args = getopt.getopt(sys.argv[1:], 'hge:i:f:c:', ["help", "gzip", "elastic", "index", "file","count"])        
    except getopt.GetoptError, e:
        print str(e)
        usage()        
        
    b_gzip = False
    es_host = '' 
    es_index = ''
    save_file = ''
    count = False
    encoding = sys.getfilesystemencoding()     
    
    for o, a in opts:
        a = a.decode(encoding)
        if o in ("-h", "--help"):
            usage()
        elif o in ("-g", "gzip"): #
            b_gzip  = True        
        elif o in ("-e", "--elastic"): 
            es_host = a
        elif o in ("-i", "--index"):
            es_index = a
        elif o in ("-f", "--file"):
            save_file = a
        elif o in ("-c", "--count"):
            count = a
        else:
            assert False, "Unhandled Option" 
            
    if not (es_host and es_index and save_file):
        usage()           
    sample = ElasticExportIndex(es_host, es_index, save_file, b_gzip, count)
    print sample
    
if __name__ == "__main__":
    main()    

print 'Exiting Main Thread...'


================================================
FILE: fdns_download.py
================================================
#coding: utf-8

import os
import sys    
import requests
import shlex
import subprocess
import hashlib
import gzip
import json
import chardet
from datetime import datetime
from bs4 import BeautifulSoup
from elasticsearch import Elasticsearch, helpers


class Elastic(object):
        
    def __init__(self, elastic=None, url='http://192.168.1.121:9200'):        
      self.elastic = elastic or Elasticsearch(url)
      
    def exists(self, index):
        return  self.elastic.indices.exists(index)      
           
    def create(self, index, body=None):                  
        self.elastic.indices.create(index, body=body, ignore=[400])   
    
    def insert(self, index, type, datas, id=None):    
        actions = []           
        for data in datas:
            actions.append({
                    "_index": index,
                    "_type": type,
                    "_id": id,
                    "_source": data
            })
        try:      
            helpers.bulk(self.elastic, actions) 
        except Exception, e:
            print ("{0} start bulk error: {1}".format(datetime.now(), e))         

class Fdns(object):
        
    def __init__(self):        
        self.website = "https://scans.io/study/sonar.fdns"       
        self.downurl, self.code = self.get_download_link()                     
        
    def get_download_link(self):       
        resp = requests.get(self.website)
        if resp.status_code != 200:
            raise Exception("download page failed: %d" % resp.status_code)
        soup = BeautifulSoup(resp.text, 'lxml')             
        try:
            table = soup.findAll(name='table', attrs={'class':'table table-condensed'})[0]
            tr = table.findAll(name='tr')[-1]              
            href = tr.findAll(name='a')[0]['href']
            code = tr.findAll(name='code')[0].text
        except Exception, msg:
            raise Exception("get download link failed: {}".format(msg))               
        return (href, code.lower())    
        
    @staticmethod
    def sha1(filepath, block_size=64*1024):
        try:
            with open(filepath, 'rb') as fd:
                sha1obj = hashlib.sha1()
                while True:
                    data = fd.read(block_size)
                    if not data:
                        break
                    sha1obj.update(data)
                retsha1 = sha1obj.hexdigest()
                return retsha1
        except IOError:
            raise Exception('Invalid file path: {}'.format(filepath))
        
    @staticmethod
    def make_dirs(dirpath, default='./sample'):
        try:       
            if not os.path.exists(dirpath):
                os.makedirs(dirpath)
            return dirpath.rstrip(os.sep)
        except Exception, e:
            if not os.path.exists(default):
                os.makedirs(default)
            return dirpath.rstrip(os.sep) 
        
    @staticmethod
    def download_file(downurl, localdir=None, showlog=False):    
        localdir = Fdns.make_dirs(localdir)
        command_line = "wget -c -t100 -P {0} {1}".format(localdir, downurl)
        tmp_cmdline = shlex.split(command_line)
        try:
            proc = subprocess.Popen(args=tmp_cmdline,
                                    stderr=subprocess.STDOUT,
                                    stdout=subprocess.PIPE,
                                    bufsize=0)
        except IOError:
             raise EnvironmentError(1, "wget is not installed or could "
                                      "not be found in system path") 
        while showlog and proc.poll() is None:
            for streamline in iter(proc.stdout.readline, ''):
                sys.stdout.write(streamline)
        proc.communicate()        
        return proc.returncode      
        
    @staticmethod
    def gzip_extract(gzpath, dstpath, block_size=64*1024):
        try:
            with gzip.open(gzpath, 'rb') as fr, open(dstpath, 'wb') as fw:
                while True:
                    data = fr.read(block_size)
                    if not data:
                        break
                    fw.write(data)
        except IOError:
            raise Exception('Invalid file path: {}'.format(gzpath))                    
        
    @staticmethod
    def import_elastic(elastic, gzfile, step=500, sep=','):        
        with gzip.open(gzfile) as fd_file_down:          
            lines = []
            type = 'json'                                                   
            index = os.path.basename(gzfile)                 
            elastic.create(index=index, body={ 
                'mappings': {
                    'json': {
                        'properties' : {
                            'domain' : {
                                'type' : 'string',
                                'index': 'not_analyzed'
                            },
                            'record_type' : {
                                'type' : 'string',
                                'index': 'not_analyzed'
                            },
                            'record_value' : {
                                'type' : 'string',
                                'index': 'not_analyzed'
                            }                                                  
                        }                
                    }                    
                }
            })
                               
            for line in fd_file_down:
                success = False
                if not success:
                    try:
                        fields = line.strip().split(sep)
                        data = dict(domain=fields[0], record_type=fields[1], record_value=fields[2])                                                                          
                        lines.append(json.dumps(data))
                        success = True
                    except Exception:
                        pass
                if not success:
                    try:
                        encoding = chardet.detect(line).get('encoding')
                        line = line.decode(encoding)
                        fields = line.strip().split(sep)
                        data = dict(domain=fields[0], record_type=fields[1], record_value=fields[2])                                                                          
                        lines.append(json.dumps(data))
                        success = True                      
                    except Exception:
                        pass                      
                if len(lines) >= step:
                    elastic.insert(index, type, lines)
                    lines[:] = []
            if len(lines):
                elastic.insert(index, type, lines)                  
                      
def main():    
    fdns = Fdns() 
    localdir = "./sample"   
    downfilename = os.path.basename(fdns.downurl)
    
    print ("{0} start download file: {1}".format(datetime.now(), downfilename))
    while True:
        returncode = fdns.download_file(fdns.downurl, localdir=localdir, showlog=False)
        if returncode == 0:
            break
    
    print ("{0} start sha1 file: {1}".format(datetime.now(), downfilename))    
    downfilepath = localdir + os.path.sep + downfilename
    if fdns.sha1(downfilepath) != fdns.code:
        print "downfile failed: {}".format(fdns.downurl)
        sys.exit()
        
    print ("{0} start import file: {1}".format(datetime.now(), downfilename)) 
    elastic = Elastic()
    if not elastic.exists(downfilename):
        fdns.import_elastic(elastic, downfilepath)    
    
if __name__ == "__main__":   
    main()



================================================
FILE: json2es.py
================================================
#coding: utf-8
import os
import sys
import json
import gzip
import threading
import Queue
import getopt
from datetime import datetime
from elasticsearch import Elasticsearch, helpers
import sys 
reload(sys) 
sys.setdefaultencoding('utf8')

class logRecode(object):
    def __init__(self):
        self.begin_time = datetime.now()
        self.end_time = datetime.now()
        self.total = 0
        self.success = 0 
        self.failed = 0   

class ImportElasticIndex(logRecode):    
    def __init__(self, es_host, es_index, src_file, 
                 b_gzip=False, scroll_size=1000, queue_size=10000):
        logRecode.__init__(self)        
        self.elastic = Elasticsearch(es_host)
        self.es_index = es_index  
        self.doc_type = 'json' 
        self.src_file = src_file
        self.b_gzip = b_gzip
        self.scroll_size = scroll_size     
        self.queue = Queue.Queue(queue_size)
        self.fquit = object()
        self.start()      
        
    def is_gzip(self, filename):
        try:
            fp = open(filename, 'rb') 
            if os.path.getsize(filename) >=2:
                magic = self.fileobj.read(2)
                if magic == '\037\213':
                    return True                    
        finally:
            fp.close()                              
        
    def read_file(self):        
        try:
            if self.b_gzip:
                fp = gzip.open(self.src_file, 'rb')
            else:
                fp = open(self.src_file, 'rb')
            for record in fp: 
                try:               
                    record = record.strip()
                    self.queue.put(record) 
                    self.total += 1             
                except Exception, e:
                    print e                               
        finally:
            fp.close()
    
    def write_elastic(self):
        actions = []                
        while True:
            record = self.queue.get()                  
            if record == self.fquit:
                self.queue.put(self.fquit)
                if len(actions):
                    try:                            
                        success, failed = helpers.bulk(self.elastic, actions, True)
                        self.success += success
                        self.failed += failed
                        actions[:] = [] 
                    except Exception, e:
                        print e                                        
                break                     
            actions.append({
                   "_index": self.es_index,
                    "_type": self.doc_type,
                    "_id": None,
                    "_source": record                            
            })                        
            if len(actions) >= self.scroll_size:
                try:
                    success, failed = helpers.bulk(self.elastic, actions, True)
                    self.success += success
                    self.failed += failed 
                    actions[:] = [] 
                except Exception, e:                  
                    print e            
            self.queue.task_done()     
        self.end_time = datetime.now()
            
    def __str__(self):
        return json.dumps({
            "begin_time": str(self.begin_time),
            "end_time:": str(self.end_time),
            "total:": self.total,
            "success:": self.success,
            "failed:":  self.failed}, indent=4)        
    
    def start(self):
        if not os.path.exists(os.path.abspath(self.src_file)):
            print "Error: Source file %s not exists." % self.src_file
            sys.exit(-1)                               
        self.elastic.indices.create(self.es_index, 
                                    body=json.loads(open('mapping.json').read()),
                                    ignore=[400, 404])        
        read_thread  = threading.Thread(target=self.read_file)
        read_thread.start()           
         
        write_thread = threading.Thread(target=self.write_elastic)
        write_thread.start()
        
        read_thread.join()
        self.queue.join()
        self.queue.put(self.fquit)
        write_thread.join()                 
        
def usage():      
    print "Example:"
    print "%s -e 192.168.1.121:9200 -i indexname -f result.json -g" %sys.argv[0]    
    sys.exit(1)
    
def main():    
    if len(sys.argv[1:]) < 1:
        usage()
    try: 
        opts, args = getopt.getopt(sys.argv[1:], 'hge:i:f:', 
                                   ["help", "gzip", "elastic", "index", "file"])        
    except getopt.GetoptError, e:
        print str(e)
        usage()        
        
    b_gzip = False
    es_host = '' 
    es_index = ''
    src_file = ''  
    encoding = sys.getfilesystemencoding()     
    
    for o, a in opts:
        a = a.decode(encoding)
        if o in ("-h", "--help"):
            usage()
        elif o in ("-g", "gzip"):
            b_gzip  = True        
        elif o in ("-e", "--elastic"): 
            es_host = a
        elif o in ("-i", "--index"):
            es_index = a
        elif o in ("-f", "--file"):
            src_file = a
        else:
            assert False, "Unhandled Option" 
            
    if not (es_host and es_index and src_file):
        usage()           
    sample = ImportElasticIndex(es_host, es_index, src_file, b_gzip)
    print sample
    
if __name__ == "__main__":
    main()