<s> import os <EOL> import json <EOL> from framework . dependency_management . dependency_resolver import BaseComponent <EOL> from framework . dependency_management . interfaces import PluginOutputInterface <EOL> from framework . db . target_manager import target_required <EOL> from framework . lib . exceptions import InvalidParameterType <EOL> from framework . db import models <EOL> from framework . utils import FileOperations <EOL> class POutputDB ( BaseComponent , PluginOutputInterface ) : <EOL> COMPONENT_NAME = "<STR_LIT>" <EOL> def __init__ ( self ) : <EOL> self . register_in_service_locator ( ) <EOL> self . config = self . get_component ( "<STR_LIT>" ) <EOL> self . plugin_handler = self . get_component ( "<STR_LIT>" ) <EOL> self . reporter = self . get_component ( "<STR_LIT>" ) <EOL> self . target = self . get_component ( "<STR_LIT:target>" ) <EOL> self . db_config = self . get_component ( "<STR_LIT>" ) <EOL> self . timer = self . get_component ( "<STR_LIT>" ) <EOL> self . db = self . get_component ( "<STR_LIT>" ) <EOL> def PluginOutputExists ( self , plugin_key , target_id ) : <EOL> count = self . db . session . query ( models . PluginOutput ) . filter_by ( target_id = target_id , plugin_key = plugin_key ) . count ( ) <EOL> return ( count > <NUM_LIT:0> ) <EOL> def PluginCountOutput ( self ) : <EOL> complete_count = self . db . session . query ( models . PluginOutput ) . count ( ) <EOL> left_count = self . db . session . query ( models . Work ) . count ( ) <EOL> results = { '<STR_LIT>' : complete_count , '<STR_LIT>' : left_count } <EOL> return results <EOL> def DeriveHTMLOutput ( self , plugin_output ) : <EOL> Content = '<STR_LIT>' <EOL> for item in plugin_output : <EOL> Content += getattr ( <EOL> self . reporter , <EOL> item [ "<STR_LIT:type>" ] ) ( ** item [ "<STR_LIT>" ] ) <EOL> return ( Content ) <EOL> @ target_required <EOL> def DeriveOutputDict ( self , obj , target_id = None ) : <EOL> if target_id : <EOL> self . target . SetTarget ( target_id ) <EOL> if obj : <EOL> pdict = dict ( obj . __dict__ ) <EOL> pdict . pop ( "<STR_LIT>" , None ) <EOL> pdict . pop ( "<STR_LIT>" ) <EOL> if pdict . get ( "<STR_LIT>" , None ) : <EOL> pdict [ "<STR_LIT>" ] = self . DeriveHTMLOutput ( <EOL> json . loads ( pdict [ "<STR_LIT>" ] ) ) <EOL> pdict [ "<STR_LIT>" ] = obj . start_time . strftime ( <EOL> self . db_config . Get ( "<STR_LIT>" ) ) <EOL> pdict [ "<STR_LIT>" ] = obj . end_time . strftime ( <EOL> self . db_config . Get ( "<STR_LIT>" ) ) <EOL> pdict [ "<STR_LIT>" ] = self . timer . get_time_as_str ( obj . run_time ) <EOL> return pdict <EOL> @ target_required <EOL> def DeriveOutputDicts ( self , obj_list , target_id = None ) : <EOL> if target_id : <EOL> self . target . SetTarget ( target_id ) <EOL> dict_list = [ ] <EOL> for obj in obj_list : <EOL> dict_list . append ( self . DeriveOutputDict ( obj , target_id = target_id ) ) <EOL> return ( dict_list ) <EOL> def GenerateQueryUsingSession ( self , filter_data , target_id , for_delete = False ) : <EOL> query = self . db . session . query ( models . PluginOutput ) . filter_by ( target_id = target_id ) <EOL> if filter_data . get ( "<STR_LIT>" , None ) : <EOL> query . filter_by ( target_id = filter_data [ "<STR_LIT>" ] ) <EOL> if filter_data . get ( "<STR_LIT>" , None ) : <EOL> if isinstance ( filter_data . get ( "<STR_LIT>" ) , ( str , unicode ) ) : <EOL> query = query . filter_by ( plugin_key = filter_data [ "<STR_LIT>" ] ) <EOL> if isinstance ( filter_data . get ( "<STR_LIT>" ) , list ) : <EOL> query = query . filter ( models . PluginOutput . plugin_key . in_ ( <EOL> filter_data [ "<STR_LIT>" ] ) ) <EOL> if filter_data . get ( "<STR_LIT>" , None ) : <EOL> if isinstance ( filter_data . get ( "<STR_LIT>" ) , ( str , unicode ) ) : <EOL> query = query . filter_by ( plugin_type = filter_data [ "<STR_LIT>" ] ) <EOL> if isinstance ( filter_data . get ( "<STR_LIT>" ) , list ) : <EOL> query = query . filter ( models . PluginOutput . plugin_type . in_ ( <EOL> filter_data [ "<STR_LIT>" ] ) ) <EOL> if filter_data . get ( "<STR_LIT>" , None ) : <EOL> if isinstance ( filter_data . get ( "<STR_LIT>" ) , ( str , unicode ) ) : <EOL> query = query . filter_by ( <EOL> plugin_group = filter_data [ "<STR_LIT>" ] ) <EOL> if isinstance ( filter_data . get ( "<STR_LIT>" ) , list ) : <EOL> query = query . filter ( models . PluginOutput . plugin_group . in_ ( <EOL> filter_data [ "<STR_LIT>" ] ) ) <EOL> if filter_data . get ( "<STR_LIT>" , None ) : <EOL> if isinstance ( filter_data . get ( "<STR_LIT>" ) , ( str , unicode ) ) : <EOL> query = query . filter_by ( plugin_code = filter_data [ "<STR_LIT>" ] ) <EOL> if isinstance ( filter_data . get ( "<STR_LIT>" ) , list ) : <EOL> query = query . filter ( models . PluginOutput . plugin_code . in_ ( <EOL> filter_data [ "<STR_LIT>" ] ) ) <EOL> if filter_data . get ( "<STR_LIT:status>" , None ) : <EOL> if isinstance ( filter_data . get ( "<STR_LIT:status>" ) , ( str , unicode ) ) : <EOL> query = query . filter_by ( status = filter_data [ "<STR_LIT:status>" ] ) <EOL> if isinstance ( filter_data . get ( "<STR_LIT:status>" ) , list ) : <EOL> query = query . filter ( models . PluginOutput . status . in_ ( <EOL> filter_data [ "<STR_LIT:status>" ] ) ) <EOL> try : <EOL> if filter_data . get ( "<STR_LIT>" , None ) : <EOL> if isinstance ( filter_data . get ( "<STR_LIT>" ) , ( str , unicode ) ) : <EOL> query = query . filter_by ( <EOL> user_rank = int ( filter_data [ "<STR_LIT>" ] ) ) <EOL> if isinstance ( filter_data . get ( "<STR_LIT>" ) , list ) : <EOL> numbers_list = [ int ( x ) for x in filter_data [ "<STR_LIT>" ] ] <EOL> query = query . filter ( models . PluginOutput . user_rank . in_ ( <EOL> numbers_list ) ) <EOL> if filter_data . get ( "<STR_LIT>" , None ) : <EOL> if isinstance ( filter_data . get ( "<STR_LIT>" ) , ( str , unicode ) ) : <EOL> query = query . filter_by ( <EOL> owtf_rank = int ( filter_data [ "<STR_LIT>" ] ) ) <EOL> if isinstance ( filter_data . get ( "<STR_LIT>" ) , list ) : <EOL> numbers_list = [ int ( x ) for x in filter_data [ "<STR_LIT>" ] ] <EOL> query = query . filter ( models . PluginOutput . owtf_rank . in_ ( <EOL> numbers_list ) ) <EOL> except ValueError : <EOL> raise InvalidParameterType ( <EOL> "<STR_LIT>" ) <EOL> if not for_delete : <EOL> query = query . order_by ( models . PluginOutput . plugin_key . asc ( ) ) <EOL> return query <EOL> @ target_required <EOL> def GetAll ( self , filter_data = None , target_id = None ) : <EOL> if not filter_data : <EOL> filter_data = { } <EOL> self . target . SetTarget ( target_id ) <EOL> query = self . GenerateQueryUsingSession ( filter_data , target_id ) <EOL> results = query . all ( ) <EOL> return ( self . DeriveOutputDicts ( results , target_id = target_id ) ) <EOL> @ target_required <EOL> def GetUnique ( self , target_id = None ) : <EOL> """<STR_LIT>""" <EOL> unique_data = { <EOL> "<STR_LIT>" : [ i [ <NUM_LIT:0> ] for i in self . db . session . query ( <EOL> models . PluginOutput . plugin_type ) . filter_by ( <EOL> target_id = target_id ) . distinct ( ) . all ( ) ] , <EOL> "<STR_LIT>" : [ i [ <NUM_LIT:0> ] for i in self . db . session . query ( <EOL> models . PluginOutput . plugin_group ) . filter_by ( <EOL> target_id = target_id ) . distinct ( ) . all ( ) ] , <EOL> "<STR_LIT:status>" : [ i [ <NUM_LIT:0> ] for i in self . db . session . query ( <EOL> models . PluginOutput . status ) . filter_by ( <EOL> target_id = target_id ) . distinct ( ) . all ( ) ] , <EOL> "<STR_LIT>" : [ i [ <NUM_LIT:0> ] for i in self . db . session . query ( <EOL> models . PluginOutput . user_rank ) . filter_by ( <EOL> target_id = target_id ) . distinct ( ) . all ( ) ] , <EOL> "<STR_LIT>" : [ i [ <NUM_LIT:0> ] for i in self . db . session . query ( <EOL> models . PluginOutput . owtf_rank ) . filter_by ( <EOL> target_id = target_id ) . distinct ( ) . all ( ) ] , <EOL> } <EOL> return ( unique_data ) <EOL> @ target_required <EOL> def DeleteAll ( self , filter_data , target_id = None ) : <EOL> """<STR_LIT>""" <EOL> query = self . GenerateQueryUsingSession ( <EOL> filter_data , <EOL> target_id , <EOL> for_delete = True ) <EOL> for plugin in query . all ( ) : <EOL> if plugin . output_path : <EOL> output_path = os . path . join ( <EOL> self . config . GetOutputDirForTargets ( ) , <EOL> plugin . output_path ) <EOL> if os . path . exists ( output_path ) : <EOL> FileOperations . rm_tree ( output_path ) <EOL> results = query . delete ( ) <EOL> self . db . session . commit ( ) <EOL> @ target_required <EOL> def Update ( self , plugin_group , plugin_type , plugin_code , patch_data , target_id = None ) : <EOL> query = self . GenerateQueryUsingSession ( { <EOL> "<STR_LIT>" : plugin_group , <EOL> "<STR_LIT>" : plugin_type , <EOL> "<STR_LIT>" : plugin_code , } , <EOL> target_id ) <EOL> obj = query . first ( ) <EOL> if obj : <EOL> try : <EOL> if patch_data . get ( "<STR_LIT>" , None ) : <EOL> if isinstance ( patch_data [ "<STR_LIT>" ] , list ) : <EOL> patch_data [ "<STR_LIT>" ] = patch_data [ "<STR_LIT>" ] [ <NUM_LIT:0> ] <EOL> obj . user_rank = int ( patch_data [ "<STR_LIT>" ] ) <EOL> if patch_data . get ( "<STR_LIT>" , None ) : <EOL> if isinstance ( patch_data [ "<STR_LIT>" ] , list ) : <EOL> patch_data [ "<STR_LIT>" ] = patch_data [ "<STR_LIT>" ] [ <NUM_LIT:0> ] <EOL> obj . user_notes = patch_data [ "<STR_LIT>" ] <EOL> self . db . session . merge ( obj ) <EOL> self . db . session . commit ( ) <EOL> except ValueError : <EOL> raise InvalidParameterType ( "<STR_LIT>" ) <EOL> def PluginAlreadyRun ( self , PluginInfo , target_id = None ) : <EOL> plugin_output_count = self . db . session . query ( models . PluginOutput ) . filter_by ( <EOL> target_id = target_id , <EOL> plugin_code = PluginInfo [ "<STR_LIT:code>" ] , <EOL> plugin_type = PluginInfo [ "<STR_LIT:type>" ] , <EOL> plugin_group = PluginInfo [ "<STR_LIT>" ] ) . count ( ) <EOL> return plugin_output_count > <NUM_LIT:0> <EOL> @ target_required <EOL> def SavePluginOutput ( self , plugin , output , target_id = None ) : <EOL> """<STR_LIT>""" <EOL> self . db . session . merge ( models . PluginOutput ( <EOL> plugin_key = plugin [ "<STR_LIT:key>" ] , <EOL> plugin_code = plugin [ "<STR_LIT:code>" ] , <EOL> plugin_group = plugin [ "<STR_LIT>" ] , <EOL> plugin_type = plugin [ "<STR_LIT:type>" ] , <EOL> output = json . dumps ( output ) , <EOL> start_time = plugin [ "<STR_LIT:start>" ] , <EOL> end_time = plugin [ "<STR_LIT:end>" ] , <EOL> status = plugin [ "<STR_LIT:status>" ] , <EOL> target_id = target_id , <EOL> output_path = ( plugin [ "<STR_LIT>" ] if os . path . exists ( <EOL> self . plugin_handler . GetPluginOutputDir ( plugin ) ) else None ) , <EOL> owtf_rank = plugin [ '<STR_LIT>' ] ) <EOL> ) <EOL> self . db . session . commit ( ) <EOL> @ target_required <EOL> def SavePartialPluginOutput ( self , plugin , output , message , target_id = None ) : <EOL> self . db . session . merge ( models . PluginOutput ( <EOL> plugin_key = plugin [ "<STR_LIT:key>" ] , <EOL> plugin_code = plugin [ "<STR_LIT:code>" ] , <EOL> plugin_group = plugin [ "<STR_LIT>" ] , <EOL> plugin_type = plugin [ "<STR_LIT:type>" ] , <EOL> output = json . dumps ( output ) , <EOL> error = message , <EOL> start_time = plugin [ "<STR_LIT:start>" ] , <EOL> end_time = plugin [ "<STR_LIT:end>" ] , <EOL> status = plugin [ "<STR_LIT:status>" ] , <EOL> target_id = target_id , <EOL> output_path = ( plugin [ "<STR_LIT>" ] if os . path . exists ( <EOL> self . plugin_handler . GetPluginOutputDir ( plugin ) ) else None ) , <EOL> owtf_rank = plugin [ '<STR_LIT>' ] ) <EOL> ) <EOL> self . db . session . commit ( ) </s>
