feat: Added attribute_whitelist and attribute_blacklist

fczbkk · Oct 17, 2018 · 6cea39f · 6cea39f
1 parent 649bc14
commit 6cea39f
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 12 deletions.
diff --git a/build/css-selector-generator.js b/build/css-selector-generator.js
@@ -6,7 +6,9 @@
     CssSelectorGenerator.prototype.default_options = {
       selectors: ['id', 'class', 'tag', 'nthchild'],
       prefix_tag: false,
-      log: false
+      log: false,
+      attribute_blacklist: [],
+      attribute_whitelist: []
     };
 
     function CssSelectorGenerator(options) {
@@ -107,14 +109,21 @@
     };
 
     CssSelectorGenerator.prototype.getAttributeSelectors = function(element) {
-      var attribute, blacklist, k, len, ref, ref1, result;
+      var a, attr, blacklist, k, l, len, len1, ref, ref1, ref2, result, whitelist;
       result = [];
-      blacklist = ['id', 'class'];
+      whitelist = this.options.attribute_whitelist;
+      for (k = 0, len = whitelist.length; k < len; k++) {
+        attr = whitelist[k];
+        if (element.hasAttribute(attr)) {
+          result.push("[" + attr + "=" + (this.sanitizeItem(element.getAttribute(attr))) + "]");
+        }
+      }
+      blacklist = this.options.attribute_blacklist.concat(['id', 'class']);
       ref = element.attributes;
-      for (k = 0, len = ref.length; k < len; k++) {
-        attribute = ref[k];
-        if (ref1 = attribute.nodeName, indexOf.call(blacklist, ref1) < 0) {
-          result.push("[" + attribute.nodeName + "=" + attribute.nodeValue + "]");
+      for (l = 0, len1 = ref.length; l < len1; l++) {
+        a = ref[l];
+        if (!((ref1 = a.nodeName, indexOf.call(blacklist, ref1) >= 0) || (ref2 = a.nodeName, indexOf.call(whitelist, ref2) >= 0))) {
+          result.push("[" + a.nodeName + "=" + (this.sanitizeItem(a.nodeValue)) + "]");
         }
       }
       return result;

diff --git a/src/css-selector-generator.coffee b/src/css-selector-generator.coffee
@@ -3,7 +3,9 @@ class CssSelectorGenerator
   default_options:
     # choose from 'tag', 'id', 'class', 'nthchild', 'attribute'
     selectors: ['id', 'class', 'tag', 'nthchild'],
-    prefix_tag: false, log: false
+    prefix_tag: false, log: false,
+    attribute_blacklist: [],
+    attribute_whitelist: []
 
   constructor: (options = {}) ->
     @options = {}
@@ -82,10 +84,14 @@ class CssSelectorGenerator
 
   getAttributeSelectors: (element) ->
     result = []
-    blacklist = ['id', 'class']
-    for attribute in element.attributes
-      unless attribute.nodeName in blacklist
-        result.push "[#{attribute.nodeName}=#{attribute.nodeValue}]"
+    whitelist = @options.attribute_whitelist
+    for attr in whitelist
+      if element.hasAttribute attr
+        result.push "[#{attr}=#{@sanitizeItem element.getAttribute(attr)}]"
+    blacklist = @options.attribute_blacklist.concat(['id', 'class'])
+    for a in element.attributes
+      unless a.nodeName in blacklist or a.nodeName in whitelist
+        result.push "[#{a.nodeName}=#{@sanitizeItem a.nodeValue}]"
     result
 
   getNthChildSelector: (element) ->

diff --git a/test/src/css-selector-generator.spec.coffee b/test/src/css-selector-generator.spec.coffee
@@ -341,6 +341,45 @@ describe 'CSS Selector Generator', ->
         expect(x.getSelector elm).
                  toEqual '[rel=bbb]'
 
+      it 'should get attribute selectors for an element ignoring blacklist', ->
+        x.setOptions attribute_blacklist: [ 'href' ]
+        elm = root.querySelector '#linkZero'
+        result = x.getAttributeSelectors elm
+        expect(result).not.toContain '[href=linkThree]'
+        expect(result).toContain '[target=someTarget]'
+        expect(result).toContain '[rel=someRel]'
+        expect(result).not.toContain '[id=linkZero]'
+        expect(result.length).toBe 2
+        expect(x.getClassSelectors root).toEqual []
+
+      it 'should get attribute selectors prioritizing whitelist', ->
+        x.setOptions attribute_whitelist: [ 'rel' ]
+        elm = root.querySelector '#linkZero'
+        result = x.getAttributeSelectors elm
+        expect(result[0]).toEqual '[rel=someRel]'
+        expect(x.getClassSelectors root).toEqual []
+
+      it 'should sanitize attribute values', ->
+        x.setOptions selectors: ['attribute']
+        x.setOptions attribute_whitelist: ['href']
+        x.setOptions log:true
+        href = "https://www.google.co.in/" + \
+               "search?rlz=1C5CHFA_enIN755IN755&ei=wwPHW7OrOdv0rQHR-" + \
+               "KmABA&q=using+set+coffeescript&oq=using+set+coffeescript" + \
+               "&gs_l=psy-ab.3...2980.7444.0.7719.22." + \
+               "22.0.0.0.0.219.2518.0j15j2.17.0....0...1c.1.64.psy-ab..5" + \
+               ".15.2263...0" + \
+               "j0i131k1j0i67k1j0i131i67k1j0i22i30k1j33i22i29i30k1j33i21" + \
+               "k1j33i160k1." + \
+               "0.8WGhjkzVME4"
+        root.innerHTML = '<div data-id="a1"><a id="aaa" href=' + \
+                         "#{href}" + ' rel="bbb">link1</a></div>' + \
+                         '<div><a href="aaa" rel="aaa">link2</a></div>'
+
+        elm = root.querySelector '#aaa'
+        expect(x.getSelector elm).
+                 toEqual "[href=#{x.sanitizeItem href}]"
+
     describe 'n-th child', ->
 
       it 'should get n-th child selector for an element', ->