From 57c96f5adf5551fe44419a53d842f7f74f54267c Mon Sep 17 00:00:00 2001
From: Beatriz Uezu <beatriz.uezu@gmail.com>
Date: Thu, 17 Oct 2019 00:27:12 -0300
Subject: [PATCH] Add validation for CNH

---
 tests/test_all.py          |  2 ++
 tests/test_cnh.py          | 42 ++++++++++++++++++++++++
 validate_docbr/CNH.py      | 67 ++++++++++++++++++++++++++++++++++++++
 validate_docbr/__init__.py |  1 +
 4 files changed, 112 insertions(+)
 create mode 100644 tests/test_cnh.py
 create mode 100644 validate_docbr/CNH.py

diff --git a/tests/test_all.py b/tests/test_all.py
index 51d497a..1bcae0f 100644
--- a/tests/test_all.py
+++ b/tests/test_all.py
@@ -1,6 +1,7 @@
 import unittest
 # Classes
 import tests.test_cpf
+import tests.test_cnh
 import tests.test_cns
 import tests.test_cnpj
 
@@ -10,6 +11,7 @@ def suite():
     test_suite = unittest.TestSuite()
 
     test_suite.addTests(loader.loadTestsFromModule(tests.test_cpf))
+    test_suite.addTests(loader.loadTestsFromModule(tests.test_cnh))
     test_suite.addTests(loader.loadTestsFromModule(tests.test_cns))
     test_suite.addTests(loader.loadTestsFromModule(tests.test_cnpj))
 
diff --git a/tests/test_cnh.py b/tests/test_cnh.py
new file mode 100644
index 0000000..a7a4d20
--- /dev/null
+++ b/tests/test_cnh.py
@@ -0,0 +1,42 @@
+import unittest
+import validate_docbr as docbr
+
+
+class TestCnh(unittest.TestCase):
+    """Testar a classe CNH."""
+
+    def setUp(self):
+        """Inicia novo objeto em todo os testes."""
+        self.cnh = docbr.CNH()
+
+    def test_generate_validate(self):
+        """Verifica os métodos de geração e validação de documento."""
+        # generate_list
+        cnhs = (
+            self.cnh.generate_list(1)
+            + self.cnh.generate_list(1, mask=True)
+            + self.cnh.generate_list(1, mask=True, repeat=True)
+        )
+        self.assertIsInstance(cnhs, list)
+        self.assertTrue(len(cnhs) == 3)
+
+        # validate_list
+        cnhs_validates = self.cnh.validate_list(cnhs)
+        self.assertTrue(sum(cnhs_validates) == 3)
+
+    def test_mask(self):
+        """Verifica se o método mask funciona corretamente."""
+        masked_cpf = self.cnh.mask('11122233344')
+        self.assertEqual(masked_cpf, '111 222 333 44')
+
+    def test_special_case(self):
+        """ Verifica os casos especiais de CNH """
+        cases = [
+            ('00000000000', False),
+            ('78623161668', False),
+            ('0123 456 789 10', False),
+            ('65821310502', True),
+            ('658 213 105 02', True),
+        ]
+        for cnh, is_valid in cases:
+            self.assertEqual(self.cnh.validate(cnh), is_valid)
diff --git a/validate_docbr/CNH.py b/validate_docbr/CNH.py
new file mode 100644
index 0000000..43a9212
--- /dev/null
+++ b/validate_docbr/CNH.py
@@ -0,0 +1,67 @@
+from .BaseDoc import BaseDoc
+from random import sample
+from typing import List
+
+
+class CNH(BaseDoc):
+    """Classe referente ao Carteira Nacional de Habilitação (CNH)."""
+
+    def __init__(self):
+        self.digits = list(range(10))
+        self.dsc = 0
+
+    def validate(self, doc: str = '') -> bool:
+        """Validar CNH."""
+        doc = self._only_digits(doc)
+
+        if len(doc) != 11 or self._is_repeated_digits(doc):
+            return False
+
+        first_digit = self._generate_first_digit(doc)
+        second_digit = self._generate_second_digit(doc)
+
+        return first_digit == doc[9] and second_digit == doc[10]
+
+    def generate(self, mask: bool = False) -> str:
+        """Gerar CNH."""
+        cnh = [str(sample(self.digits, 1)[0]) for i in range(9)]
+        cnh.append(self._generate_first_digit(cnh))
+        cnh.append(self._generate_second_digit(cnh))
+
+        cnh = ''.join(cnh)
+        return self.mask(cnh) if mask else cnh
+
+    def mask(self, doc: str = '') -> str:
+        """Coloca a máscara de CNH na variável doc."""
+        return f'{doc[:3]} {doc[3:6]} {doc[6:9]} {doc[9:]}'
+
+    def _generate_first_digit(self, doc: list) -> tuple:
+        """Gerar o primeiro dígito verificador da CNH."""
+        sum = 0
+
+        for i in range(9, 0, -1):
+            sum += int(doc[9 - i]) * i
+
+        first_value = sum % 11
+        if first_value >= 10:
+            first_value, self.dsc = 0, 2
+        return str(first_value)
+
+    def _generate_second_digit(self, doc: list) -> str:
+        """Gerar o segundo dígito verificador da CNH."""
+        sum = 0
+
+        for i in range(1, 10):
+            sum += int(doc[i-1]) * i
+
+        rest = sum % 11
+
+        second_value = rest - self.dsc
+        if rest >= 10:
+            second_value = 0
+        return str(second_value)
+
+    def _is_repeated_digits(self, doc: List[str]) -> bool:
+        """Verifica se é uma CNH contém com números repetidos.
+        Exemplo: 11111111111"""
+        return len(set(doc)) == 1
diff --git a/validate_docbr/__init__.py b/validate_docbr/__init__.py
index 1f84556..b8dff56 100644
--- a/validate_docbr/__init__.py
+++ b/validate_docbr/__init__.py
@@ -1,4 +1,5 @@
 from .BaseDoc import BaseDoc
 from .CPF import CPF
 from .CNPJ import CNPJ
+from .CNH import CNH
 from .CNS import CNS