UrlCategorise 0.1.3 → 0.1.6

data/lib/url_categorise/client.rb

@@ -3,6 +3,7 @@ require 'set'
 module UrlCategorise
   class Client < ApiPattern::Client
     include ::UrlCategorise::Constants
+    include ActiveAttr::Model
 
     def self.compatible_api_version
       'v2'
@@ -12,16 +13,35 @@ module UrlCategorise
       'v2 2025-08-23'
     end
 
-    attr_reader :host_urls, :hosts, :cache_dir, :force_download, :dns_servers, :metadata, :request_timeout,
-                :dataset_processor, :dataset_categories
+    attribute :host_urls, default: -> { DEFAULT_HOST_URLS }
+    attribute :cache_dir
+    attribute :force_download, type: Boolean, default: false
+    attribute :dns_servers, default: ['1.1.1.1', '1.0.0.1']
+    attribute :request_timeout, type: Integer, default: 10
+    attribute :iab_compliance_enabled, type: Boolean, default: false
+    attribute :iab_version, default: :v3
+    attribute :auto_load_datasets, type: Boolean, default: false
+    attribute :smart_categorization_enabled, type: Boolean, default: false
+    attribute :smart_rules, default: -> { {} }
+
+    attr_reader :hosts, :metadata, :dataset_processor, :dataset_categories
+
+    def initialize(**kwargs)
+      # Extract dataset_config for later use
+      dataset_config = kwargs.fetch(:dataset_config, {})
+      
+      # Set ActiveAttr attributes - preserve explicitly passed values including nil
+      self.host_urls = kwargs.key?(:host_urls) ? kwargs[:host_urls] : DEFAULT_HOST_URLS
+      self.cache_dir = kwargs[:cache_dir] # will be nil if not provided or explicitly nil
+      self.force_download = kwargs.key?(:force_download) ? kwargs[:force_download] : false
+      self.dns_servers = kwargs.key?(:dns_servers) ? kwargs[:dns_servers] : ['1.1.1.1', '1.0.0.1']
+      self.request_timeout = kwargs.key?(:request_timeout) ? kwargs[:request_timeout] : 10
+      self.iab_compliance_enabled = kwargs.key?(:iab_compliance) ? kwargs[:iab_compliance] : false
+      self.iab_version = kwargs.key?(:iab_version) ? kwargs[:iab_version] : :v3
+      self.auto_load_datasets = kwargs.key?(:auto_load_datasets) ? kwargs[:auto_load_datasets] : false
+      self.smart_categorization_enabled = kwargs.key?(:smart_categorization) ? kwargs[:smart_categorization] : false
+      self.smart_rules = initialize_smart_rules(kwargs.key?(:smart_rules) ? kwargs[:smart_rules] : {})
 
-    def initialize(host_urls: DEFAULT_HOST_URLS, cache_dir: nil, force_download: false,
-                   dns_servers: ['1.1.1.1', '1.0.0.1'], request_timeout: 10, dataset_config: {})
-      @host_urls = host_urls
-      @cache_dir = cache_dir
-      @force_download = force_download
-      @dns_servers = dns_servers
-      @request_timeout = request_timeout
       @metadata = {}
       @dataset_categories = Set.new # Track which categories come from datasets
 
@@ -29,30 +49,48 @@ module UrlCategorise
       @dataset_processor = initialize_dataset_processor(dataset_config) unless dataset_config.empty?
 
       @hosts = fetch_and_build_host_lists
+
+      # Auto-load datasets from constants if enabled
+      load_datasets_from_constants if auto_load_datasets && @dataset_processor
     end
 
     def categorise(url)
       host = (URI.parse(url).host || url).downcase
       host = host.gsub('www.', '')
 
-      @hosts.keys.select do |category|
+      categories = @hosts.keys.select do |category|
         @hosts[category].any? do |blocked_host|
           host == blocked_host || host.end_with?(".#{blocked_host}")
         end
       end
+
+      # Apply smart categorization if enabled
+      categories = apply_smart_categorization(url, categories) if smart_categorization_enabled
+
+      if iab_compliance_enabled
+        IabCompliance.get_iab_categories(categories, iab_version)
+      else
+        categories
+      end
     end
 
     def categorise_ip(ip_address)
-      @hosts.keys.select do |category|
+      categories = @hosts.keys.select do |category|
         @hosts[category].include?(ip_address)
       end
+
+      if iab_compliance_enabled
+        IabCompliance.get_iab_categories(categories, iab_version)
+      else
+        categories
+      end
     end
 
     def resolve_and_categorise(domain)
       categories = categorise(domain)
 
       begin
-        resolver = Resolv::DNS.new(nameserver: @dns_servers)
+        resolver = Resolv::DNS.new(nameserver: dns_servers)
         ip_addresses = resolver.getaddresses(domain).map(&:to_s)
 
         ip_addresses.each do |ip|
@@ -79,6 +117,62 @@ module UrlCategorise
       hash_size_in_mb(@hosts)
     end
 
+    def size_of_dataset_data
+      dataset_hosts = {}
+      @dataset_categories.each do |category|
+        dataset_hosts[category] = @hosts[category] || []
+      end
+      hash_size_in_mb(dataset_hosts)
+    end
+
+    def size_of_blocklist_data
+      blocklist_hosts = {}
+      @hosts.each do |category, domains|
+        blocklist_hosts[category] = domains unless @dataset_categories.include?(category)
+      end
+      hash_size_in_mb(blocklist_hosts)
+    end
+
+    def size_of_data_bytes
+      hash_size_in_bytes(@hosts)
+    end
+
+    def size_of_dataset_data_bytes
+      dataset_hosts = {}
+      @dataset_categories.each do |category|
+        dataset_hosts[category] = @hosts[category] || []
+      end
+      hash_size_in_bytes(dataset_hosts)
+    end
+
+    def size_of_blocklist_data_bytes
+      blocklist_hosts = {}
+      @hosts.each do |category, domains|
+        blocklist_hosts[category] = domains unless @dataset_categories.include?(category)
+      end
+      hash_size_in_bytes(blocklist_hosts)
+    end
+
+    def count_of_dataset_hosts
+      @dataset_categories.map do |category|
+        @hosts[category]&.size || 0
+      end.sum
+    end
+
+    def count_of_dataset_categories
+      @dataset_categories.size
+    end
+
+    def iab_compliant?
+      iab_compliance_enabled
+    end
+
+    def get_iab_mapping(category)
+      return nil unless iab_compliance_enabled
+
+      IabCompliance.map_category_to_iab(category, iab_version)
+    end
+
     def check_all_lists
       puts 'Checking all lists in constants...'
 
@@ -86,7 +180,7 @@ module UrlCategorise
       missing_categories = []
       successful_lists = {}
 
-      @host_urls.each do |category, urls|
+      (host_urls || {}).each do |category, urls|
         puts "\nChecking category: #{category}"
 
         if urls.empty?
@@ -114,7 +208,7 @@ module UrlCategorise
           print "  🔍 Testing #{url}... "
 
           begin
-            response = HTTParty.head(url, timeout: @request_timeout, follow_redirects: true)
+            response = HTTParty.head(url, timeout: request_timeout, follow_redirects: true)
 
             case response.code
             when 200
@@ -166,7 +260,7 @@ module UrlCategorise
       puts '=' * 80
 
       puts "\n📊 SUMMARY:"
-      total_categories = @host_urls.keys.length
+      total_categories = (host_urls || {}).keys.length
       categories_with_issues = unreachable_lists.keys.length + missing_categories.length
       categories_healthy = total_categories - categories_with_issues
 
@@ -266,16 +360,163 @@ module UrlCategorise
         @hosts[category].concat(domains).uniq!
       end
 
+      # Reload datasets from constants if auto-loading is enabled
+      load_datasets_from_constants if auto_load_datasets && @dataset_processor
+
       self
     end
 
+    def export_hosts_files(output_path = nil)
+      export_dir = output_path || (cache_dir ? File.join(cache_dir, 'exports', 'hosts') : File.join(Dir.pwd, 'exports', 'hosts'))
+      
+      FileUtils.mkdir_p(export_dir) unless Dir.exist?(export_dir)
+      
+      exported_files = {}
+      
+      @hosts.each do |category, domains|
+        next if domains.empty?
+        
+        filename = "#{category}.hosts"
+        file_path = File.join(export_dir, filename)
+        
+        File.open(file_path, 'w') do |file|
+          file.puts "# #{category.to_s.gsub('_', ' ').split.map(&:capitalize).join(' ')} - Generated by UrlCategorise"
+          file.puts "# Generated at: #{Time.now}"
+          file.puts "# Total entries: #{domains.length}"
+          file.puts ""
+          
+          domains.sort.each do |domain|
+            file.puts "0.0.0.0 #{domain}"
+          end
+        end
+        
+        exported_files[category] = {
+          path: file_path,
+          filename: filename,
+          count: domains.length
+        }
+      end
+      
+      # Create summary file
+      summary_path = File.join(export_dir, '_export_summary.txt')
+      File.open(summary_path, 'w') do |file|
+        file.puts "UrlCategorise Hosts Export Summary"
+        file.puts "=================================="
+        file.puts "Generated at: #{Time.now}"
+        file.puts "Export directory: #{export_dir}"
+        file.puts "Total categories: #{exported_files.keys.length}"
+        file.puts "Total domains: #{@hosts.values.map(&:length).sum}"
+        file.puts ""
+        file.puts "Files created:"
+        
+        exported_files.each do |category, info|
+          file.puts "  #{info[:filename]} - #{info[:count]} domains"
+        end
+      end
+      
+      exported_files[:_summary] = {
+        path: summary_path,
+        total_categories: exported_files.keys.length,
+        total_domains: @hosts.values.map(&:length).sum,
+        export_directory: export_dir
+      }
+      
+      exported_files
+    end
+
+    def export_csv_data(output_path = nil)
+      require 'csv'
+      
+      export_dir = output_path || (cache_dir ? File.join(cache_dir, 'exports', 'csv') : File.join(Dir.pwd, 'exports', 'csv'))
+      
+      FileUtils.mkdir_p(export_dir) unless Dir.exist?(export_dir)
+      
+      filename = "url_categorise_data_export_#{Time.now.strftime('%Y%m%d_%H%M%S')}.csv"
+      file_path = File.join(export_dir, filename)
+      
+      CSV.open(file_path, 'w', headers: true) do |csv|
+        # Add headers
+        csv << [
+          'domain',
+          'category', 
+          'source_type',
+          'is_dataset_category',
+          'iab_category_v2',
+          'iab_category_v3',
+          'export_timestamp',
+          'smart_categorization_enabled'
+        ]
+        
+        # Export all host/category data
+        @hosts.each do |category, domains|
+          domains.each do |domain|
+            source_type = @dataset_categories.include?(category) ? 'dataset' : 'blocklist'
+            is_dataset_category = @dataset_categories.include?(category)
+            
+            # Get IAB mappings if compliance is enabled
+            iab_v2 = nil
+            iab_v3 = nil
+            if iab_compliance_enabled
+              iab_v2 = IabCompliance.map_category_to_iab(category, :v2)
+              iab_v3 = IabCompliance.map_category_to_iab(category, :v3)
+            end
+            
+            csv << [
+              domain,
+              category,
+              source_type,
+              is_dataset_category,
+              iab_v2,
+              iab_v3,
+              Time.now.iso8601,
+              smart_categorization_enabled
+            ]
+          end
+        end
+      end
+      
+      # Create metadata file
+      metadata_path = File.join(export_dir, "#{File.basename(filename, '.csv')}_metadata.json")
+      metadata = {
+        export_info: {
+          timestamp: Time.now.iso8601,
+          filename: filename,
+          file_path: file_path,
+          metadata_path: metadata_path
+        },
+        client_settings: {
+          iab_compliance_enabled: iab_compliance_enabled,
+          iab_version: iab_version,
+          smart_categorization_enabled: smart_categorization_enabled,
+          auto_load_datasets: auto_load_datasets
+        },
+        data_summary: {
+          total_domains: @hosts.values.map(&:length).sum,
+          total_categories: @hosts.keys.length,
+          dataset_categories_count: @dataset_categories.size,
+          blocklist_categories_count: @hosts.keys.length - @dataset_categories.size,
+          categories: @hosts.keys.sort.map(&:to_s)
+        },
+        dataset_metadata: dataset_metadata
+      }
+      
+      File.write(metadata_path, JSON.pretty_generate(metadata))
+      
+      {
+        csv_file: file_path,
+        metadata_file: metadata_path,
+        summary: metadata[:data_summary],
+        export_directory: export_dir
+      }
+    end
+
     private
 
     def initialize_dataset_processor(config)
       processor_config = {
-        download_path: config[:download_path] || @cache_dir&.+(File::SEPARATOR + 'downloads'),
-        cache_path: config[:cache_path] || @cache_dir&.+(File::SEPARATOR + 'datasets'),
-        timeout: config[:timeout] || @request_timeout,
+        download_path: config[:download_path] || cache_dir&.+(File::SEPARATOR + 'downloads'),
+        cache_path: config[:cache_path] || cache_dir&.+(File::SEPARATOR + 'datasets'),
+        timeout: config[:timeout] || request_timeout,
         enable_kaggle: config.fetch(:enable_kaggle, true) # Default to true for backwards compatibility
       }
 
@@ -298,6 +539,7 @@ module UrlCategorise
 
     def integrate_dataset(dataset, category_mappings)
       return dataset unless @dataset_processor
+      return nil unless dataset # Handle nil datasets gracefully
 
       categorized_data = @dataset_processor.integrate_dataset_into_categorization(dataset, category_mappings)
 
@@ -324,21 +566,104 @@ module UrlCategorise
       dataset
     end
 
+    def load_datasets_from_constants
+      return unless defined?(CATEGORIY_DATABASES) && CATEGORIY_DATABASES.is_a?(Array)
+      return unless @dataset_processor
+
+      puts "Loading #{CATEGORIY_DATABASES.length} datasets from constants..." if ENV['DEBUG']
+      loaded_count = 0
+
+      CATEGORIY_DATABASES.each do |dataset_config|
+        begin
+          case dataset_config[:type]
+          when :kaggle
+            # Parse the kaggle path to get owner and dataset name
+            path_parts = dataset_config[:path].split('/')
+            next unless path_parts.length == 2
+
+            dataset_owner, dataset_name = path_parts
+            
+            # Check if dataset is already cached before attempting to load
+            cache_key = @dataset_processor.send(:generate_cache_key, "#{dataset_owner}/#{dataset_name}", :kaggle)
+            cache_file = File.join(@dataset_processor.cache_path, cache_key)
+            
+            if File.exist?(cache_file)
+              puts "Loading cached Kaggle dataset: #{dataset_owner}/#{dataset_name}" if ENV['DEBUG']
+              load_kaggle_dataset(dataset_owner, dataset_name, {
+                                    use_cache: true,
+                                    integrate_data: true
+                                  })
+              loaded_count += 1
+            else
+              puts "Attempting to download missing Kaggle dataset: #{dataset_owner}/#{dataset_name}" if ENV['DEBUG']
+              begin
+                load_kaggle_dataset(dataset_owner, dataset_name, {
+                                      use_cache: true,
+                                      integrate_data: true
+                                    })
+                loaded_count += 1
+              rescue Error => e
+                puts "Warning: Failed to download Kaggle dataset #{dataset_owner}/#{dataset_name}: #{e.message}" if ENV['DEBUG']
+              end
+            end
+
+          when :csv
+            # Check if CSV dataset is cached
+            cache_key = @dataset_processor.send(:generate_cache_key, dataset_config[:path], :csv)
+            cache_file = File.join(@dataset_processor.cache_path, cache_key)
+            
+            if File.exist?(cache_file)
+              puts "Loading cached CSV dataset: #{dataset_config[:path]}" if ENV['DEBUG']
+              load_csv_dataset(dataset_config[:path], {
+                                 use_cache: true,
+                                 integrate_data: true
+                               })
+              loaded_count += 1
+            else
+              puts "Attempting to download missing CSV dataset: #{dataset_config[:path]}" if ENV['DEBUG']
+              begin
+                load_csv_dataset(dataset_config[:path], {
+                                   use_cache: true,
+                                   integrate_data: true
+                                 })
+                loaded_count += 1
+              rescue Error => e
+                puts "Warning: Failed to download CSV dataset #{dataset_config[:path]}: #{e.message}" if ENV['DEBUG']
+              end
+            end
+          end
+        rescue Error => e
+          puts "Warning: Failed to load dataset #{dataset_config[:path]}: #{e.message}" if ENV['DEBUG']
+          # Continue loading other datasets even if one fails
+        rescue StandardError => e
+          puts "Warning: Unexpected error loading dataset #{dataset_config[:path]}: #{e.message}" if ENV['DEBUG']
+          # Continue loading other datasets even if one fails
+        end
+      end
+
+      puts "Finished loading datasets from constants (#{loaded_count}/#{CATEGORIY_DATABASES.length} loaded)" if ENV['DEBUG']
+    end
+
     def hash_size_in_mb(hash)
-      size = 0
+      size_bytes = hash_size_in_bytes(hash)
+      (size_bytes / ONE_MEGABYTE.to_f).round(2)
+    end
 
+    def hash_size_in_bytes(hash)
+      size = 0
       hash.each do |_key, value|
+        next unless value.is_a?(Array)
+
         size += value.join.length
       end
-
-      (size / ONE_MEGABYTE).round(2)
+      size
     end
 
     def fetch_and_build_host_lists
       @hosts = {}
 
-      host_urls.keys.each do |category|
-        @hosts[category] = build_host_data(host_urls[category])
+      (host_urls || {}).keys.each do |category|
+        @hosts[category] = build_host_data((host_urls || {})[category])
       end
 
       sub_category_values = categories_with_keys
@@ -356,6 +681,79 @@ module UrlCategorise
       @hosts
     end
 
+    def initialize_smart_rules(custom_rules)
+      custom_rules = {} if custom_rules.nil?
+      default_rules = {
+        social_media_platforms: {
+          domains: %w[reddit.com facebook.com twitter.com x.com instagram.com linkedin.com
+                      pinterest.com tiktok.com youtube.com snapchat.com discord.com],
+          remove_categories: %i[health_and_fitness forums news technology education
+                                business finance entertainment travel sports politics
+                                science music art food_and_drink shopping gaming]
+        },
+        search_engines: {
+          domains: %w[google.com bing.com yahoo.com duckduckgo.com baidu.com yandex.com],
+          remove_categories: %i[news shopping travel health_and_fitness finance technology]
+        },
+        video_platforms: {
+          domains: %w[youtube.com vimeo.com dailymotion.com twitch.tv],
+          remove_categories: %i[education news entertainment music sports gaming]
+        },
+        news_aggregators: {
+          domains: %w[reddit.com digg.com],
+          keep_primary_only: %i[social_media reddit digg]
+        }
+      }
+
+      # Merge custom rules with defaults
+      default_rules.merge(custom_rules)
+    end
+
+    def apply_smart_categorization(url, categories)
+      return categories unless smart_categorization_enabled
+
+      host = extract_host(url)
+
+      smart_rules.each do |_rule_name, rule_config|
+        if rule_config[:domains]&.any? { |domain| host == domain || host.end_with?(".#{domain}") }
+          categories = apply_rule(categories, rule_config, host, url)
+        end
+      end
+
+      categories
+    end
+
+    def apply_rule(categories, rule_config, _host, url)
+      # Rule: Remove overly broad categories for specific platforms
+      if rule_config[:remove_categories]
+        categories = categories.reject { |cat| rule_config[:remove_categories].include?(cat) }
+      end
+
+      # Rule: Keep only primary categories
+      if rule_config[:keep_primary_only]
+        primary_categories = categories & rule_config[:keep_primary_only]
+        categories = primary_categories if primary_categories.any?
+      end
+
+      # Rule: Add specific categories based on URL patterns
+      if rule_config[:add_categories_by_path]
+        rule_config[:add_categories_by_path].each do |path_pattern, additional_categories|
+          categories = (categories + additional_categories).uniq if url.match?(path_pattern)
+        end
+      end
+
+      # Rule: Remove all categories except allowed ones
+      categories &= rule_config[:allowed_categories_only] if rule_config[:allowed_categories_only]
+
+      categories
+    end
+
+    def extract_host(url)
+      (URI.parse(url).host || url).downcase.gsub('www.', '')
+    rescue URI::InvalidURIError
+      url.downcase.gsub('www.', '')
+    end
+
     def build_host_data(urls)
       all_hosts = []
 
@@ -364,11 +762,11 @@ module UrlCategorise
 
         hosts_data = nil
 
-        hosts_data = read_from_cache(url) if @cache_dir && !@force_download
+        hosts_data = read_from_cache(url) if cache_dir && !force_download
 
         if hosts_data.nil?
           hosts_data = download_and_parse_list(url)
-          save_to_cache(url, hosts_data) if @cache_dir
+          save_to_cache(url, hosts_data) if cache_dir
         end
 
         all_hosts.concat(hosts_data) if hosts_data
@@ -378,7 +776,7 @@ module UrlCategorise
     end
 
     def download_and_parse_list(url)
-      raw_data = HTTParty.get(url, timeout: @request_timeout)
+      raw_data = HTTParty.get(url, timeout: request_timeout)
       return [] if raw_data.body.nil? || raw_data.body.empty?
 
       # Store metadata
@@ -441,11 +839,11 @@ module UrlCategorise
     end
 
     def cache_file_path(url)
-      return nil unless @cache_dir
+      return nil unless cache_dir
 
-      FileUtils.mkdir_p(@cache_dir) unless Dir.exist?(@cache_dir)
+      FileUtils.mkdir_p(cache_dir) unless Dir.exist?(cache_dir)
       filename = Digest::MD5.hexdigest(url) + '.cache'
-      File.join(@cache_dir, filename)
+      File.join(cache_dir, filename)
     end
 
     def read_from_cache(url)
@@ -478,7 +876,7 @@ module UrlCategorise
     end
 
     def should_update_cache?(url, cache_data)
-      return true if @force_download
+      return true if force_download
       return true unless cache_data[:metadata]
 
       # Update if cache is older than 24 hours
@@ -487,7 +885,7 @@ module UrlCategorise
 
       # Check if remote content has changed
       begin
-        head_response = HTTParty.head(url, timeout: @request_timeout)
+        head_response = HTTParty.head(url, timeout: request_timeout)
         remote_etag = head_response.headers['etag']
         remote_last_modified = head_response.headers['last-modified']
 
@@ -508,8 +906,8 @@ module UrlCategorise
     def categories_with_keys
       keyed_categories = {}
 
-      host_urls.keys.each do |category|
-        category_values = host_urls[category].select do |url|
+      (host_urls || {}).keys.each do |category|
+        category_values = (host_urls || {})[category].select do |url|
           url.is_a?(Symbol)
         end
 

data/lib/url_categorise/dataset_processor.rb

@@ -58,11 +58,16 @@ module UrlCategorise
         return handle_existing_dataset(extracted_dir, options)
       end
 
-      # Only require credentials if we need to download fresh data
+      # If credentials not available, return nil gracefully for cache mode
       unless kaggle_credentials_available?
-        raise Error, 'Kaggle credentials required for downloading new datasets. ' \
-                     'Set KAGGLE_USERNAME/KAGGLE_KEY environment variables, provide credentials explicitly, ' \
-                     'or place kaggle.json file in ~/.kaggle/ directory.'
+        if options[:use_cache]
+          puts "Warning: Kaggle dataset '#{dataset_path}' not cached and no credentials available" if ENV['DEBUG']
+          return nil
+        else
+          raise Error, 'Kaggle credentials required for downloading new datasets. ' \
+                       'Set KAGGLE_USERNAME/KAGGLE_KEY environment variables, provide credentials explicitly, ' \
+                       'or place kaggle.json file in ~/.kaggle/ directory.'
+        end
       end
 
       # Download from Kaggle API