# HG changeset patch
# User Benoît Allard <benoit.allard@greenbone.net>
# Date 1419939019 -3600
# Node ID 3826f2701ff23d8c9c638a629897bab76a18825c
# Parent  652f59fbea3afa2d9c1cf8e1030c3d045cab7945
CPE: Add the possibility to add ourself integrally to the product tree

diff -r 652f59fbea3a -r 3826f2701ff2 farolluz/parsers/cpe.py
--- a/farolluz/parsers/cpe.py	Tue Dec 30 12:29:07 2014 +0100
+++ b/farolluz/parsers/cpe.py	Tue Dec 30 12:30:19 2014 +0100
@@ -36,6 +36,15 @@
 
 import re
 
+from ..producttree import CVRFFullProductName, CVRFRelationship
+
+def capitalize(s):
+    """ A custom version of string.capwords that split on _, and join on ' '
+    """
+    s = s.replace('\\', '')
+    return ' '.join(c.capitalize() for c in s.split('_'))
+
+
 PCT_MAP ={'!': "%21", '"': "%22", '#': "%23", '$': "%24", '%': "%25", '&': "%26",
           "'": "%27", '(': "%28", ')': "%29", '*': "%2a", '+': "%2b", ',': "%2c",
           '/': "%2f", ':': "%3a", ';': "%3b", '<': "%3c", "=": "%3d", '>': "%3e",
@@ -282,6 +291,74 @@
             elif idx == 12:
                 self.other = v
 
+    def addToDoc(self, document, finalProduct=True):
+        """ Add the CPE value as full producttree in the document
+        If finalProduct is false, only the elements leading to the product
+        will be added.
+        """
+        ptree = document._producttree
+        if ptree is None:
+            ptree = document.createProductTree()
+
+        def next_prodid():
+            """ A handy function to generate the next available productid """
+            prods = document._producttree._products
+            if len(prods) > 0:
+                last_prodid = prods[-1]._productid
+                numlen = 0
+                while last_prodid[- (numlen + 1)] in "0123456789":
+                    numlen += 1
+                if numlen != 0:
+                    return last_prodid[:-numlen] + str(int(last_prodid[-numlen:]) + 1)
+            return document.getDocId() + '-P0'
+
+        # Create the main product tree
+        tree = []
+        for value, valtype in [(self.vendor, 'Vendor'),
+                            (self.product, 'Product Name'),
+                            (self.version, 'Product Version'),
+                            (self.update, 'Patch Level'),
+                            (self.language, 'Language'),
+                            (self.target_hw, 'Architecture')]:
+            if value.value is not None:
+                tree.append((valtype, capitalize(value.value)))
+
+        # Import it
+        last_branch = ptree.importTree(tree)
+        # Add a product there
+        if self.target_sw.value is None:
+            if not finalProduct:
+                return last_branch
+            product = CVRFFullProductName(next_prodid(), str(self), last_branch, self.bind_to_fs())
+            ptree.addProduct(product)
+            return product
+        else:
+            product = CVRFFullProductName(next_prodid(), str(self), last_branch)
+            ptree.addProduct(product)
+
+        # We do have a target software, we need to create a relationship !
+        os = CVRFFullProductName(next_prodid(), self.target_sw.value, ptree)
+        ptree.addProduct(os)
+
+        rel = CVRFRelationship(product._productid, 'Installed On', os._productid)
+        ptree.addRelationship(rel)
+        if not finalProduct:
+            return rel
+
+        final_prod = CVRFFullProductName(next_prodid(), ptree.getNameOfRelationship(rel), rel, self.bind_to_fs())
+        ptree.addProduct(final_prod)
+        return final_prod
+
+    def __str__(self):
+        res = []
+        if self.product.value:
+            res.append(capitalize(self.product.value))
+        if self.version.value:
+            res.append(capitalize(self.version.value))
+        if not res:
+            return capitalize(self.vendor.value)
+        return ' '.join(res)
+
 def parse(s):
     cpe = CPE()
     if s[:5] == 'cpe:/':
diff -r 652f59fbea3a -r 3826f2701ff2 farolluz/parsers/cve.py
--- a/farolluz/parsers/cve.py	Tue Dec 30 12:29:07 2014 +0100
+++ b/farolluz/parsers/cve.py	Tue Dec 30 12:30:19 2014 +0100
@@ -32,6 +32,7 @@
 
 import xml.etree.ElementTree as ET
 
+from .cpe import parse as parseCPE
 from .xml import parseDate
 
 from .. import __version__
@@ -150,13 +151,7 @@
     for i, cpe in enumerate(xml.findall(
                                 '/'.join([UN('vuln', 'vulnerable-software-list'),
                                           UN('vuln', 'product')]))):
-        if doc._producttree is None:
-            doc.createProductTree()
-        try:
-            prod = doc._producttree.getProductForCPE(cpe.text)
-        except KeyError:
-            prod = CVRFFullProductName('%s-P%d' % (vulnid, i), cpe.text, doc._producttree, cpe.text)
-            doc._producttree.addProduct(prod)
+        prod = parseCPE(cpe.text).addToDoc(doc)
         vulnerable_products.append(prod)
 
     if vulnerable_products: