#!/usr/bin/env ruby
# encoding: UTF-8

# Authors:
# Krzysztof Kotowicz - @kkotowicz - http://blog.kotowicz.net

require 'rubygems'
require 'json'
require 'securerandom'

class ChromeExtensionToolkit
    @ext = ''
    @manifest
    BCG_SCRIPT = 1
    BCG_PAGE = 2

    def initialize(extpath)
        raise ArgumentError, "Empty extension path" unless extpath
        raise ArgumentError, "Invalid extension path" unless File.directory?(extpath)
        @ext = extpath
    end

    def reload_manifest()
        f = get_file('manifest.json')
        manifest = File.read(f).sub("\xef\xbb\xbf", '')
        manifest = JSON.parse(manifest)
        set_manifest(manifest)
        return manifest
    end

    def get_manifest()
        if @manifest
            return @manifest
        end

        return reload_manifest()
    end

    def set_manifest(manifest)
        @manifest = manifest
    end

    def save_manifest()
        save_file('manifest.json', JSON.pretty_generate(@manifest))
    end

    def save_file(file, contents)
        File.open(get_file(file), 'w') {|f| f.write contents }
    end

    def get_file(file)
        return File.join(@ext, file)
    end

    def assert_not_app()
        manifest = get_manifest()
        raise RuntimeError,  "Apps are not supported, only regular Chrome extensions" if manifest['app']
    end

    def inject_script(payload)
        assert_not_app()
        assert_background_page('injector_bg') # add page to extensions that don't have one
        bcg_file = get_file(get_background_page())

        if File.exist?(bcg_file)
            bcg = File.read(bcg_file)
        else
            bcg = ""
        end

        if not payload
            return bcg
        end

        if bcg_file.end_with? ".js" # js file, just prepend payload
            return payload + ";" + bcg
        end

        name = SecureRandom.hex + '.js'
        save_file(name, payload)
        return bcg.sub(/(\<head\>|\Z)/i, "\\1\n<script src=\"/#{name}\"></script>")
    end

    def assert_background_page(default)
        manifest = get_manifest()
        if not manifest['manifest_version'].nil? and manifest['manifest_version'] >= 2
            if manifest['background'].nil?
                manifest['background'] = {}
            end

            if manifest['background']['page']
                return BCG_PAGE
            end

            if not manifest['background']['scripts'].nil?
                manifest['background']['scripts'].unshift(default + '.js')
                set_manifest(manifest)
                return BCG_SCRIPT
            else

                manifest['background']['scripts'] = [default + '.js']
                set_manifest(manifest)
                return BCG_SCRIPT
            end
        end
        if not manifest['background_page']
            manifest['background_page'] = default + '.html'
        end
        set_manifest(manifest)
        return BCG_PAGE
    end

    def add_permissions(perms)
        manifest = get_manifest()
        manifest['permissions'] += perms
        manifest['permissions'].uniq!
        set_manifest(manifest)
    end

    def get_background_page()
        manifest = get_manifest()
        if manifest['background'] and manifest['background']['page']
            return manifest['background']['page']
        end

        if manifest['background'] and manifest['background']['scripts']
            return manifest['background']['scripts'][0]
        end

        if manifest['background_page']
            return manifest['background_page']
        end
        raise RuntimeError, "No background page present"
    end
end